src

Go monorepo.
git clone git://code.dwrz.net/src
Log | Files | Refs

index.go (2347B)

      1 // Copyright 2026 The Go Authors. All rights reserved.
      2 // Use of this source code is governed by a BSD-style
      3 // license that can be found in the LICENSE file.
      4 
      5 package graph
      6 
      7 import (
      8 	"fmt"
      9 	"iter"
     10 	"slices"
     11 )
     12 
     13 // An Index is an immutable, bijective map between [0, N) and an ordered list of keys.
     14 type Index[Key comparable] struct {
     15 	// There are three Index representations:
     16 	//
     17 	// - If identN > 0, an identity map of [0, identN).
     18 	// - If index == nil, a sorted integer index in values.
     19 	// - Otherwise, a full index in values and index.
     20 
     21 	identN int
     22 	values []Key
     23 	index  map[Key]int
     24 }
     25 
     26 // NewIndex returns an index for the specified list of values.
     27 func NewIndex[Key comparable](values iter.Seq[Key]) *Index[Key] {
     28 	vs := slices.Collect(values)
     29 	if len(vs) == 0 {
     30 		return new(Index[Key])
     31 	}
     32 
     33 	// Fast path: a naturally sorted list needs no index. (Sadly, there's no way
     34 	// to ask "is Key ordered?")
     35 	if vi, ok := any(vs).([]int); ok && slices.IsSorted(vi) {
     36 		return &Index[Key]{values: vs}
     37 	}
     38 
     39 	index := make(map[Key]int, len(vs))
     40 	for i, v := range vs {
     41 		index[v] = i
     42 	}
     43 	return &Index[Key]{values: vs, index: index}
     44 }
     45 
     46 // NewIdentityIndex returns an index that maps [0, n) to [0, n).
     47 func NewIdentityIndex(n int) *Index[int] {
     48 	if n < 0 {
     49 		panic("n < 0")
     50 	}
     51 	// If n == 0, this is actually a "sorted integer index", but it doesn't
     52 	// matter because everything is out of bounds either way.
     53 	return &Index[int]{identN: n}
     54 }
     55 
     56 // Value maps an index to a key.
     57 func (ix *Index[Key]) Value(index int) Key {
     58 	if ix.identN > 0 {
     59 		if index < 0 || index >= ix.identN {
     60 			panic(fmt.Sprintf("index %d out of range [0, %d)", index, ix.identN))
     61 		}
     62 		return any(index).(Key)
     63 	}
     64 	if index < 0 || index >= len(ix.values) {
     65 		panic(fmt.Sprintf("index %d out of range [0, %d)", index, ix.identN))
     66 	}
     67 	return ix.values[index]
     68 }
     69 
     70 // Index maps a key to an index.
     71 func (ix *Index[Key]) Index(key Key) int {
     72 	if key, ok := any(key).(int); ok {
     73 		// Integer-only optimizations.
     74 		switch {
     75 		case ix.identN > 0:
     76 			// Identity.
     77 			if 0 <= key && key < ix.identN {
     78 				return key
     79 			}
     80 			goto oob
     81 
     82 		case ix.index == nil:
     83 			// Sorted integers.
     84 			if i, ok := slices.BinarySearch(any(ix.values).([]int), key); ok {
     85 				return i
     86 			}
     87 			goto oob
     88 		}
     89 	}
     90 	if i, ok := ix.index[key]; ok {
     91 		return i
     92 	}
     93 
     94 oob:
     95 	panic(fmt.Sprintf("key %v not in index", key))
     96 }