collection.gno
13.22 Kb · 527 lines
1// Package collection provides a generic collection implementation with support for
2// multiple indexes, including unique indexes and case-insensitive indexes.
3// It is designed to be used with any type and allows efficient lookups using
4// different fields or computed values.
5//
6// Example usage:
7//
8// // Define a data type
9// type User struct {
10// Name string
11// Email string
12// Age int
13// Username string
14// Tags []string
15// }
16//
17// // Create a new collection
18// c := collection.New()
19//
20// // Add indexes with different options
21// c.AddIndex("name", func(v any) string {
22// return v.(*User).Name
23// }, UniqueIndex)
24//
25// c.AddIndex("email", func(v any) string {
26// return v.(*User).Email
27// }, UniqueIndex|CaseInsensitiveIndex)
28//
29// c.AddIndex("age", func(v any) string {
30// return strconv.Itoa(v.(*User).Age)
31// }, DefaultIndex) // Non-unique index
32//
33// c.AddIndex("username", func(v any) string {
34// return v.(*User).Username
35// }, UniqueIndex|SparseIndex) // Allow empty usernames
36//
37// // For tags, we index all tags for the user
38// c.AddIndex("tag", func(v any) []string {
39// return v.(*User).Tags
40// }, DefaultIndex) // Non-unique to allow multiple users with same tag
41//
42// // Store an object
43// id := c.Set(&User{
44// Name: "Alice",
45// Email: "[email protected]",
46// Age: 30,
47// Tags: []string{"admin", "moderator"}, // User can have multiple tags
48// })
49//
50// // Retrieve by any index
51// entry := c.GetFirst("email", "[email protected]")
52// adminUsers := c.GetAll("tag", "admin") // Find all users with admin tag
53// modUsers := c.GetAll("tag", "moderator") // Find all users with moderator tag
54//
55// Index options can be combined using the bitwise OR operator.
56// Available options:
57// - DefaultIndex: Regular index with no special behavior
58// - UniqueIndex: Ensures values are unique within the index
59// - CaseInsensitiveIndex: Makes string comparisons case-insensitive
60// - SparseIndex: Skips indexing empty values (nil or empty string)
61//
62// Example: UniqueIndex|CaseInsensitiveIndex for a case-insensitive unique index
63//
64// # Versioning
65//
66// This is the B+ tree successor to [gno.land/p/nym-moultest123/collection/v0] (which is
67// backed by an AVL tree): a bump to v3 because the backing data structure — and
68// thus the on-chain storage layout — changed. The exported API is unchanged from
69// v2; only the persisted representation differs, so it is a compatibility (not a
70// source) change. A B+ tree packs many entries per persisted node, so each index
71// entry costs materially less storage and gas than the AVL backing; prefer v3
72// when the collection is part of persisted realm state.
73//
74// Two operational caveats inherited from the in-place-mutating B+ tree backing:
75//
76// - do NOT mutate the collection (Set/Update/Delete) from inside an index
77// iteration callback — the AVL backing's copy-on-write tolerated it, this one
78// does not;
79// - do NOT copy a non-zero Collection by value — the copies would share live
80// tree nodes while their state diverges.
81package collection
82
83import (
84 "errors"
85 "strings"
86
87 "gno.land/p/nt/bptree/v0"
88 "gno.land/p/nt/seqid/v0"
89)
90
91// New creates a new Collection instance with an initialized ID index.
92// The ID index is a special unique index that is always present and
93// serves as the primary key for all objects in the collection.
94func New() *Collection {
95 c := &Collection{
96 indexes: make(map[string]*Index),
97 idGen: seqid.ID(0),
98 }
99 // Initialize _id index
100 c.indexes[IDIndex] = &Index{
101 options: UniqueIndex,
102 tree: bptree.NewBPTree32(),
103 }
104 return c
105}
106
107// Collection represents a collection of objects with multiple indexes
108type Collection struct {
109 indexes map[string]*Index
110 idGen seqid.ID
111}
112
113const (
114 // IDIndex is the reserved name for the primary key index
115 IDIndex = "_id"
116)
117
118// IndexOption represents configuration options for an index using bit flags
119type IndexOption uint64
120
121const (
122 // DefaultIndex is a basic index with no special options
123 DefaultIndex IndexOption = 0
124
125 // UniqueIndex ensures no duplicate values are allowed
126 UniqueIndex IndexOption = 1 << iota
127
128 // CaseInsensitiveIndex automatically converts string values to lowercase
129 CaseInsensitiveIndex
130
131 // SparseIndex only indexes non-empty values
132 SparseIndex
133)
134
135// Index represents an index with its configuration and data.
136// The index function can return either:
137// - string: for single-value indexes
138// - []string: for multi-value indexes where one object can be indexed under multiple keys
139//
140// The backing tree stores either a single ID or []string for multiple IDs per key.
141type Index struct {
142 fn any
143 options IndexOption
144 tree bptree.ITree
145}
146
147// AddIndex adds a new index to the collection with the specified options
148//
149// Parameters:
150// - name: the unique name of the index (e.g., "tags")
151// - indexFn: a function that extracts either a string or []string from an object
152// - options: bit flags for index configuration (e.g., UniqueIndex)
153func (c *Collection) AddIndex(name string, indexFn any, options IndexOption) {
154 if name == IDIndex {
155 panic("_id is a reserved index name")
156 }
157 c.indexes[name] = &Index{
158 fn: indexFn,
159 options: options,
160 tree: bptree.NewBPTree32(),
161 }
162}
163
164// storeIndex handles how we store an ID in the index tree
165func (idx *Index) store(key string, idStr string) {
166 stored := idx.tree.Get(key)
167 if stored == nil {
168 // First entry for this key
169 idx.tree.Set(key, idStr)
170 return
171 }
172
173 // Handle existing entries
174 switch existing := stored.(type) {
175 case string:
176 if existing == idStr {
177 return // Already stored
178 }
179 // Convert to array
180 idx.tree.Set(key, []string{existing, idStr})
181 case []string:
182 // Check if ID already exists
183 for _, id := range existing {
184 if id == idStr {
185 return
186 }
187 }
188 // Append new ID
189 idx.tree.Set(key, append(existing, idStr))
190 }
191}
192
193// removeIndex handles how we remove an ID from the index tree
194func (idx *Index) remove(key string, idStr string) {
195 stored := idx.tree.Get(key)
196 if stored == nil {
197 return
198 }
199
200 switch existing := stored.(type) {
201 case string:
202 if existing == idStr {
203 idx.tree.Remove(key)
204 }
205 case []string:
206 newIds := make([]string, 0, len(existing))
207 for _, id := range existing {
208 if id != idStr {
209 newIds = append(newIds, id)
210 }
211 }
212 if len(newIds) == 0 {
213 idx.tree.Remove(key)
214 } else if len(newIds) == 1 {
215 idx.tree.Set(key, newIds[0])
216 } else {
217 idx.tree.Set(key, newIds)
218 }
219 }
220}
221
222// generateKeys extracts one or more keys from an object for a given index.
223func generateKeys(idx *Index, obj any) ([]string, bool) {
224 if obj == nil {
225 return nil, false
226 }
227
228 switch fnTyped := idx.fn.(type) {
229 case func(any) string:
230 // Single-value index
231 key := fnTyped(obj)
232 return []string{key}, true
233 case func(any) []string:
234 // Multi-value index
235 keys := fnTyped(obj)
236 return keys, true
237 default:
238 panic("invalid index function type")
239 }
240}
241
242// Set adds or updates an object in the collection.
243// Returns a positive ID if successful.
244// Returns 0 if:
245// - The object is nil
246// - A uniqueness constraint would be violated
247// - Index generation fails for any index
248func (c *Collection) Set(obj any) uint64 {
249 if obj == nil {
250 return 0
251 }
252
253 // Generate new ID
254 id := c.idGen.Next()
255 idStr := id.String()
256
257 // Check uniqueness constraints first
258 for name, idx := range c.indexes {
259 if name == IDIndex {
260 continue
261 }
262 keys, ok := generateKeys(idx, obj)
263 if !ok {
264 return 0
265 }
266
267 for _, key := range keys {
268 // Skip empty values for sparse indexes
269 if idx.options&SparseIndex != 0 && key == "" {
270 continue
271 }
272 if idx.options&CaseInsensitiveIndex != 0 {
273 key = strings.ToLower(key)
274 }
275 // Only check uniqueness for unique + single-value indexes
276 // (UniqueIndex is ambiguous; skipping that scenario)
277 if idx.options&UniqueIndex != 0 {
278 if existing := idx.tree.Get(key); existing != nil {
279 return 0
280 }
281 }
282 }
283 }
284
285 // Store in _id index first (the actual object)
286 c.indexes[IDIndex].tree.Set(idStr, obj)
287
288 // Store in all other indexes
289 for name, idx := range c.indexes {
290 if name == IDIndex {
291 continue
292 }
293 keys, ok := generateKeys(idx, obj)
294 if !ok {
295 // Rollback: remove from _id index
296 c.indexes[IDIndex].tree.Remove(idStr)
297 return 0
298 }
299
300 for _, key := range keys {
301 if idx.options&SparseIndex != 0 && key == "" {
302 continue
303 }
304 if idx.options&CaseInsensitiveIndex != 0 {
305 key = strings.ToLower(key)
306 }
307 idx.store(key, idStr)
308 }
309 }
310
311 return uint64(id)
312}
313
314// Get retrieves entries matching the given key in the specified index.
315// Returns an iterator over the matching entries.
316func (c *Collection) Get(indexName string, key string) EntryIterator {
317 idx, exists := c.indexes[indexName]
318 if !exists {
319 return EntryIterator{err: errors.New("index not found: " + indexName)}
320 }
321
322 if idx.options&CaseInsensitiveIndex != 0 {
323 key = strings.ToLower(key)
324 }
325
326 if indexName == IDIndex {
327 // For ID index, validate the ID format first
328 _, err := seqid.FromString(key)
329 if err != nil {
330 return EntryIterator{err: err}
331 }
332 }
333
334 return EntryIterator{
335 collection: c,
336 indexName: indexName,
337 key: key,
338 }
339}
340
341// GetFirst returns the first matching entry or nil if none found
342func (c *Collection) GetFirst(indexName, key string) *Entry {
343 iter := c.Get(indexName, key)
344 if iter.Next() {
345 return iter.Value()
346 }
347 return nil
348}
349
350// Delete removes an object by its ID and returns true if something was deleted
351func (c *Collection) Delete(id uint64) bool {
352 idStr := seqid.ID(id).String()
353
354 // Get the object first to clean up other indexes
355 obj := c.indexes[IDIndex].tree.Get(idStr)
356 if obj == nil {
357 return false
358 }
359
360 // Remove from all indexes
361 for name, idx := range c.indexes {
362 if name == IDIndex {
363 idx.tree.Remove(idStr)
364 continue
365 }
366 keys, ok := generateKeys(idx, obj)
367 if !ok {
368 continue
369 }
370 for _, key := range keys {
371 if idx.options&CaseInsensitiveIndex != 0 {
372 key = strings.ToLower(key)
373 }
374 idx.remove(key, idStr)
375 }
376 }
377 return true
378}
379
380// Update updates an existing object and returns true if successful
381// Returns true if the update was successful.
382// Returns false if:
383// - The object is nil
384// - The ID doesn't exist
385// - A uniqueness constraint would be violated
386// - Index generation fails for any index
387//
388// If the update fails, the collection remains unchanged.
389func (c *Collection) Update(id uint64, obj any) bool {
390 if obj == nil {
391 return false
392 }
393 idStr := seqid.ID(id).String()
394 oldObj := c.indexes[IDIndex].tree.Get(idStr)
395 if oldObj == nil {
396 return false
397 }
398
399 // Check unique constraints
400 for name, idx := range c.indexes {
401 if name == IDIndex {
402 continue
403 }
404
405 if idx.options&UniqueIndex != 0 {
406 newKeys, newOk := generateKeys(idx, obj)
407 _, oldOk := generateKeys(idx, oldObj)
408 if !newOk || !oldOk {
409 return false
410 }
411
412 for _, newKey := range newKeys {
413 if idx.options&CaseInsensitiveIndex != 0 {
414 newKey = strings.ToLower(newKey)
415 }
416
417 found := idx.tree.Get(newKey)
418 if found != nil {
419 if storedID, ok := found.(string); !ok || storedID != idStr {
420 return false
421 }
422 }
423 }
424 }
425 }
426
427 // Store old index entries for potential rollback
428 oldEntries := make(map[string][]string)
429 for name, idx := range c.indexes {
430 if name == IDIndex {
431 continue
432 }
433 oldKeys, ok := generateKeys(idx, oldObj)
434 if !ok {
435 continue
436 }
437 var adjusted []string
438 for _, okey := range oldKeys {
439 if idx.options&CaseInsensitiveIndex != 0 {
440 okey = strings.ToLower(okey)
441 }
442 // Remove the oldObj from the index right away
443 idx.remove(okey, idStr)
444 adjusted = append(adjusted, okey)
445 }
446 oldEntries[name] = adjusted
447 }
448
449 // Update the object in the _id index
450 c.indexes[IDIndex].tree.Set(idStr, obj)
451
452 // Add new index entries
453 for name, idx := range c.indexes {
454 if name == IDIndex {
455 continue
456 }
457 newKeys, ok := generateKeys(idx, obj)
458 if !ok {
459 // Rollback: restore old object and old index entries
460 c.indexes[IDIndex].tree.Set(idStr, oldObj)
461 for idxName, keys := range oldEntries {
462 for _, oldKey := range keys {
463 c.indexes[idxName].store(oldKey, idStr)
464 }
465 }
466 return false
467 }
468 for _, nkey := range newKeys {
469 if idx.options&CaseInsensitiveIndex != 0 {
470 nkey = strings.ToLower(nkey)
471 }
472 idx.store(nkey, idStr)
473 }
474 }
475
476 return true
477}
478
479// GetAll retrieves all entries matching the given key in the specified index.
480func (c *Collection) GetAll(indexName string, key string) []Entry {
481 idx, exists := c.indexes[indexName]
482 if !exists {
483 return nil
484 }
485
486 if idx.options&CaseInsensitiveIndex != 0 {
487 key = strings.ToLower(key)
488 }
489
490 if indexName == IDIndex {
491 if obj := idx.tree.Get(key); obj != nil {
492 return []Entry{{ID: key, Obj: obj}}
493 }
494 return nil
495 }
496
497 idData := idx.tree.Get(key)
498 if idData == nil {
499 return nil
500 }
501
502 // Handle both single and multi-value cases based on the actual data type
503 switch stored := idData.(type) {
504 case []string:
505 result := make([]Entry, 0, len(stored))
506 for _, idStr := range stored {
507 if obj := c.indexes[IDIndex].tree.Get(idStr); obj != nil {
508 result = append(result, Entry{ID: idStr, Obj: obj})
509 }
510 }
511 return result
512 case string:
513 if obj := c.indexes[IDIndex].tree.Get(stored); obj != nil {
514 return []Entry{{ID: stored, Obj: obj}}
515 }
516 }
517 return nil
518}
519
520// GetIndex returns the underlying tree for an index
521func (c *Collection) GetIndex(name string) bptree.ITree {
522 idx, exists := c.indexes[name]
523 if !exists {
524 return nil
525 }
526 return idx.tree
527}
