collection.gno
12.29 Kb ยท 509 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: "alice@example.com",
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", "alice@example.com")
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
63package collection
64
65import (
66 "errors"
67 "strings"
68
69 "gno.land/p/demo/avl"
70 "gno.land/p/demo/seqid"
71)
72
73// New creates a new Collection instance with an initialized ID index.
74// The ID index is a special unique index that is always present and
75// serves as the primary key for all objects in the collection.
76func New() *Collection {
77 c := &Collection{
78 indexes: make(map[string]*Index),
79 idGen: seqid.ID(0),
80 }
81 // Initialize _id index
82 c.indexes[IDIndex] = &Index{
83 options: UniqueIndex,
84 tree: avl.NewTree(),
85 }
86 return c
87}
88
89// Collection represents a collection of objects with multiple indexes
90type Collection struct {
91 indexes map[string]*Index
92 idGen seqid.ID
93}
94
95const (
96 // IDIndex is the reserved name for the primary key index
97 IDIndex = "_id"
98)
99
100// IndexOption represents configuration options for an index using bit flags
101type IndexOption uint64
102
103const (
104 // DefaultIndex is a basic index with no special options
105 DefaultIndex IndexOption = 0
106
107 // UniqueIndex ensures no duplicate values are allowed
108 UniqueIndex IndexOption = 1 << iota
109
110 // CaseInsensitiveIndex automatically converts string values to lowercase
111 CaseInsensitiveIndex
112
113 // SparseIndex only indexes non-empty values
114 SparseIndex
115)
116
117// Index represents an index with its configuration and data.
118// The index function can return either:
119// - string: for single-value indexes
120// - []string: for multi-value indexes where one object can be indexed under multiple keys
121//
122// The backing tree stores either a single ID or []string for multiple IDs per key.
123type Index struct {
124 fn any
125 options IndexOption
126 tree avl.ITree
127}
128
129// AddIndex adds a new index to the collection with the specified options
130//
131// Parameters:
132// - name: the unique name of the index (e.g., "tags")
133// - indexFn: a function that extracts either a string or []string from an object
134// - options: bit flags for index configuration (e.g., UniqueIndex)
135func (c *Collection) AddIndex(name string, indexFn any, options IndexOption) {
136 if name == IDIndex {
137 panic("_id is a reserved index name")
138 }
139 c.indexes[name] = &Index{
140 fn: indexFn,
141 options: options,
142 tree: avl.NewTree(),
143 }
144}
145
146// storeIndex handles how we store an ID in the index tree
147func (idx *Index) store(key string, idStr string) {
148 stored, exists := idx.tree.Get(key)
149 if !exists {
150 // First entry for this key
151 idx.tree.Set(key, idStr)
152 return
153 }
154
155 // Handle existing entries
156 switch existing := stored.(type) {
157 case string:
158 if existing == idStr {
159 return // Already stored
160 }
161 // Convert to array
162 idx.tree.Set(key, []string{existing, idStr})
163 case []string:
164 // Check if ID already exists
165 for _, id := range existing {
166 if id == idStr {
167 return
168 }
169 }
170 // Append new ID
171 idx.tree.Set(key, append(existing, idStr))
172 }
173}
174
175// removeIndex handles how we remove an ID from the index tree
176func (idx *Index) remove(key string, idStr string) {
177 stored, exists := idx.tree.Get(key)
178 if !exists {
179 return
180 }
181
182 switch existing := stored.(type) {
183 case string:
184 if existing == idStr {
185 idx.tree.Remove(key)
186 }
187 case []string:
188 newIds := make([]string, 0, len(existing))
189 for _, id := range existing {
190 if id != idStr {
191 newIds = append(newIds, id)
192 }
193 }
194 if len(newIds) == 0 {
195 idx.tree.Remove(key)
196 } else if len(newIds) == 1 {
197 idx.tree.Set(key, newIds[0])
198 } else {
199 idx.tree.Set(key, newIds)
200 }
201 }
202}
203
204// generateKeys extracts one or more keys from an object for a given index.
205func generateKeys(idx *Index, obj any) ([]string, bool) {
206 if obj == nil {
207 return nil, false
208 }
209
210 switch fnTyped := idx.fn.(type) {
211 case func(any) string:
212 // Single-value index
213 key := fnTyped(obj)
214 return []string{key}, true
215 case func(any) []string:
216 // Multi-value index
217 keys := fnTyped(obj)
218 return keys, true
219 default:
220 panic("invalid index function type")
221 }
222}
223
224// Set adds or updates an object in the collection.
225// Returns a positive ID if successful.
226// Returns 0 if:
227// - The object is nil
228// - A uniqueness constraint would be violated
229// - Index generation fails for any index
230func (c *Collection) Set(obj any) uint64 {
231 if obj == nil {
232 return 0
233 }
234
235 // Generate new ID
236 id := c.idGen.Next()
237 idStr := id.String()
238
239 // Check uniqueness constraints first
240 for name, idx := range c.indexes {
241 if name == IDIndex {
242 continue
243 }
244 keys, ok := generateKeys(idx, obj)
245 if !ok {
246 return 0
247 }
248
249 for _, key := range keys {
250 // Skip empty values for sparse indexes
251 if idx.options&SparseIndex != 0 && key == "" {
252 continue
253 }
254 if idx.options&CaseInsensitiveIndex != 0 {
255 key = strings.ToLower(key)
256 }
257 // Only check uniqueness for unique + single-value indexes
258 // (UniqueIndex is ambiguous; skipping that scenario)
259 if idx.options&UniqueIndex != 0 {
260 if existing, exists := idx.tree.Get(key); exists && existing != nil {
261 return 0
262 }
263 }
264 }
265 }
266
267 // Store in _id index first (the actual object)
268 c.indexes[IDIndex].tree.Set(idStr, obj)
269
270 // Store in all other indexes
271 for name, idx := range c.indexes {
272 if name == IDIndex {
273 continue
274 }
275 keys, ok := generateKeys(idx, obj)
276 if !ok {
277 // Rollback: remove from _id index
278 c.indexes[IDIndex].tree.Remove(idStr)
279 return 0
280 }
281
282 for _, key := range keys {
283 if idx.options&SparseIndex != 0 && key == "" {
284 continue
285 }
286 if idx.options&CaseInsensitiveIndex != 0 {
287 key = strings.ToLower(key)
288 }
289 idx.store(key, idStr)
290 }
291 }
292
293 return uint64(id)
294}
295
296// Get retrieves entries matching the given key in the specified index.
297// Returns an iterator over the matching entries.
298func (c *Collection) Get(indexName string, key string) EntryIterator {
299 idx, exists := c.indexes[indexName]
300 if !exists {
301 return EntryIterator{err: errors.New("index not found: " + indexName)}
302 }
303
304 if idx.options&CaseInsensitiveIndex != 0 {
305 key = strings.ToLower(key)
306 }
307
308 if indexName == IDIndex {
309 // For ID index, validate the ID format first
310 _, err := seqid.FromString(key)
311 if err != nil {
312 return EntryIterator{err: err}
313 }
314 }
315
316 return EntryIterator{
317 collection: c,
318 indexName: indexName,
319 key: key,
320 }
321}
322
323// GetFirst returns the first matching entry or nil if none found
324func (c *Collection) GetFirst(indexName, key string) *Entry {
325 iter := c.Get(indexName, key)
326 if iter.Next() {
327 return iter.Value()
328 }
329 return nil
330}
331
332// Delete removes an object by its ID and returns true if something was deleted
333func (c *Collection) Delete(id uint64) bool {
334 idStr := seqid.ID(id).String()
335
336 // Get the object first to clean up other indexes
337 obj, exists := c.indexes[IDIndex].tree.Get(idStr)
338 if !exists {
339 return false
340 }
341
342 // Remove from all indexes
343 for name, idx := range c.indexes {
344 if name == IDIndex {
345 idx.tree.Remove(idStr)
346 continue
347 }
348 keys, ok := generateKeys(idx, obj)
349 if !ok {
350 continue
351 }
352 for _, key := range keys {
353 if idx.options&CaseInsensitiveIndex != 0 {
354 key = strings.ToLower(key)
355 }
356 idx.remove(key, idStr)
357 }
358 }
359 return true
360}
361
362// Update updates an existing object and returns true if successful
363// Returns true if the update was successful.
364// Returns false if:
365// - The object is nil
366// - The ID doesn't exist
367// - A uniqueness constraint would be violated
368// - Index generation fails for any index
369//
370// If the update fails, the collection remains unchanged.
371func (c *Collection) Update(id uint64, obj any) bool {
372 if obj == nil {
373 return false
374 }
375 idStr := seqid.ID(id).String()
376 oldObj, exists := c.indexes[IDIndex].tree.Get(idStr)
377 if !exists {
378 return false
379 }
380
381 // Check unique constraints
382 for name, idx := range c.indexes {
383 if name == IDIndex {
384 continue
385 }
386
387 if idx.options&UniqueIndex != 0 {
388 newKeys, newOk := generateKeys(idx, obj)
389 _, oldOk := generateKeys(idx, oldObj)
390 if !newOk || !oldOk {
391 return false
392 }
393
394 for _, newKey := range newKeys {
395 if idx.options&CaseInsensitiveIndex != 0 {
396 newKey = strings.ToLower(newKey)
397 }
398
399 found, _ := idx.tree.Get(newKey)
400 if found != nil {
401 if storedID, ok := found.(string); !ok || storedID != idStr {
402 return false
403 }
404 }
405 }
406 }
407 }
408
409 // Store old index entries for potential rollback
410 oldEntries := make(map[string][]string)
411 for name, idx := range c.indexes {
412 if name == IDIndex {
413 continue
414 }
415 oldKeys, ok := generateKeys(idx, oldObj)
416 if !ok {
417 continue
418 }
419 var adjusted []string
420 for _, okey := range oldKeys {
421 if idx.options&CaseInsensitiveIndex != 0 {
422 okey = strings.ToLower(okey)
423 }
424 // Remove the oldObj from the index right away
425 idx.remove(okey, idStr)
426 adjusted = append(adjusted, okey)
427 }
428 oldEntries[name] = adjusted
429 }
430
431 // Update the object in the _id index
432 c.indexes[IDIndex].tree.Set(idStr, obj)
433
434 // Add new index entries
435 for name, idx := range c.indexes {
436 if name == IDIndex {
437 continue
438 }
439 newKeys, ok := generateKeys(idx, obj)
440 if !ok {
441 // Rollback: restore old object and old index entries
442 c.indexes[IDIndex].tree.Set(idStr, oldObj)
443 for idxName, keys := range oldEntries {
444 for _, oldKey := range keys {
445 c.indexes[idxName].store(oldKey, idStr)
446 }
447 }
448 return false
449 }
450 for _, nkey := range newKeys {
451 if idx.options&CaseInsensitiveIndex != 0 {
452 nkey = strings.ToLower(nkey)
453 }
454 idx.store(nkey, idStr)
455 }
456 }
457
458 return true
459}
460
461// GetAll retrieves all entries matching the given key in the specified index.
462func (c *Collection) GetAll(indexName string, key string) []Entry {
463 idx, exists := c.indexes[indexName]
464 if !exists {
465 return nil
466 }
467
468 if idx.options&CaseInsensitiveIndex != 0 {
469 key = strings.ToLower(key)
470 }
471
472 if indexName == IDIndex {
473 if obj, exists := idx.tree.Get(key); exists {
474 return []Entry{{ID: key, Obj: obj}}
475 }
476 return nil
477 }
478
479 idData, exists := idx.tree.Get(key)
480 if !exists {
481 return nil
482 }
483
484 // Handle both single and multi-value cases based on the actual data type
485 switch stored := idData.(type) {
486 case []string:
487 result := make([]Entry, 0, len(stored))
488 for _, idStr := range stored {
489 if obj, exists := c.indexes[IDIndex].tree.Get(idStr); exists {
490 result = append(result, Entry{ID: idStr, Obj: obj})
491 }
492 }
493 return result
494 case string:
495 if obj, exists := c.indexes[IDIndex].tree.Get(stored); exists {
496 return []Entry{{ID: stored, Obj: obj}}
497 }
498 }
499 return nil
500}
501
502// GetIndex returns the underlying tree for an index
503func (c *Collection) GetIndex(name string) avl.ITree {
504 idx, exists := c.indexes[name]
505 if !exists {
506 return nil
507 }
508 return idx.tree
509}
