From d73023e82c65fd6131bd905fbbe6b8ab42e3f7de Mon Sep 17 00:00:00 2001
From: moul <94029+moul@users.noreply.github.com>
Date: Tue, 10 Dec 2024 19:46:43 +0100
Subject: [PATCH] chore: fixup

Signed-off-by: moul <94029+moul@users.noreply.github.com>
---
 examples/gno.land/p/demo/avl/index/index.gno  | 62 ++++++++++++++++---
 .../gno.land/p/demo/avl/index/index_test.gno  |  1 -
 2 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/examples/gno.land/p/demo/avl/index/index.gno b/examples/gno.land/p/demo/avl/index/index.gno
index e70aed813f8..36895bddc6c 100644
--- a/examples/gno.land/p/demo/avl/index/index.gno
+++ b/examples/gno.land/p/demo/avl/index/index.gno
@@ -74,13 +74,17 @@ func (it *IndexedTree) GetByIndex(indexName string, indexKey string) []interface
 	}
 
 	var results []interface{}
-	index.tree.Iterate(indexKey, indexKey, func(key string, value interface{}) bool {
-		primaryKey := value.(string) // The index stores primary keys as values
-		if primaryValue, exists := it.primary.Get(primaryKey); exists {
-			results = append(results, primaryValue)
+
+	// Get all primary keys that match the index key
+	value, exists := index.tree.Get(indexKey)
+	if exists {
+		primaryKeys := value.([]string)
+		for _, primaryKey := range primaryKeys {
+			if primaryValue, exists := it.primary.Get(primaryKey); exists {
+				results = append(results, primaryValue)
+			}
 		}
-		return false // Continue iteration
-	})
+	}
 
 	return results
 }
@@ -90,13 +94,55 @@ func (it *IndexedTree) GetByIndex(indexName string, indexKey string) []interface
 func (it *IndexedTree) addToIndexes(primaryKey string, value interface{}) {
 	for _, index := range it.indexes {
 		indexKey := index.extract(value)
-		index.tree.Set(indexKey, primaryKey)
+
+		// Get existing keys or create new slice
+		var keys []string
+		if existing, exists := index.tree.Get(indexKey); exists {
+			keys = existing.([]string)
+		}
+		keys = append(keys, primaryKey)
+
+		index.tree.Set(indexKey, keys)
 	}
 }
 
 func (it *IndexedTree) removeFromIndexes(primaryKey string, value interface{}) {
 	for _, index := range it.indexes {
 		indexKey := index.extract(value)
-		index.tree.Remove(indexKey)
+
+		// Get existing keys
+		if existing, exists := index.tree.Get(indexKey); exists {
+			keys := existing.([]string)
+			// Remove the primary key from the slice
+			newKeys := make([]string, 0)
+			for _, k := range keys {
+				if k != primaryKey {
+					newKeys = append(newKeys, k)
+				}
+			}
+			// If there are still keys, update the index, otherwise remove it
+			if len(newKeys) > 0 {
+				index.tree.Set(indexKey, newKeys)
+			} else {
+				index.tree.Remove(indexKey)
+			}
+		}
+	}
+}
+
+// Add these methods to IndexedTree to satisfy avl.TreeInterface
+func (it *IndexedTree) Get(key string) (interface{}, bool) {
+	return it.primary.Get(key)
+}
+
+func (it *IndexedTree) Iterate(start, end string, cb func(key string, value interface{}) bool) bool {
+	return it.primary.Iterate(start, end, cb)
+}
+
+// Add this method to get access to an index as an avl.TreeInterface
+func (it *IndexedTree) GetIndexTree(name string) avl.TreeInterface {
+	if idx, exists := it.indexes[name]; exists {
+		return idx.tree
 	}
+	return nil
 }
diff --git a/examples/gno.land/p/demo/avl/index/index_test.gno b/examples/gno.land/p/demo/avl/index/index_test.gno
index 2ec19cd20cf..bf1a76b8108 100644
--- a/examples/gno.land/p/demo/avl/index/index_test.gno
+++ b/examples/gno.land/p/demo/avl/index/index_test.gno
@@ -33,7 +33,6 @@ func TestIndexedTree(t *testing.T) {
 
 	// Query by name index
 	aliceResults := it.GetByIndex("name", "Alice")
-	println(aliceResults)
 	if len(aliceResults) != 1 {
 		t.Error("Expected 1 result for Alice")
 	}