arbo: handle the case when delete when only 1 element in the tree

cmd/cli/vocdonicli.go

@@ -31,6 +31,7 @@ func (c *Config) Load(filepath string) error {
 	data, err := os.ReadFile(filepath)
 	if err != nil {
 		if errors.Is(err, os.ErrNotExist) {
+			c.LastAccountUsed = -1
 			return nil
 		}
 		return err
@@ -94,7 +95,7 @@ func NewVocdoniCLI(configFile, host string) (*vocdoniCLI, error) {
 	if err != nil {
 		return nil, err
 	}
-	if len(cfg.Accounts)-1 >= cfg.LastAccountUsed {
+	if len(cfg.Accounts)-1 >= cfg.LastAccountUsed && cfg.LastAccountUsed >= 0 {
 		log.Infof("using account %d", cfg.LastAccountUsed)
 		if err := api.SetAccount(cfg.Accounts[cfg.LastAccountUsed].PrivKey.String()); err != nil {
 			return nil, err

tree/arbo/navigate.go

@@ -11,7 +11,6 @@ import (
 // down goes down to the leaf recursively
 func (t *Tree) down(rTx db.Reader, newKey, currKey []byte, siblings [][]byte, intermediates *[][]byte,
 	path []bool, currLvl int, getLeaf bool) ([]byte, []byte, [][]byte, error) {
-
 	if currLvl > t.maxLevels {
 		return nil, nil, nil, ErrMaxLevel
 	}
@@ -184,10 +183,11 @@ func ReadIntermediateChilds(b []byte) ([]byte, []byte) {
 
 func getPath(numLevels int, k []byte) []bool {
 	requiredLen := (numLevels + 7) / 8 // Calculate the ceil value of numLevels/8
+
+	// If the provided key is shorter than expected, extend it with zero bytes
 	if len(k) < requiredLen {
-		// The provided key is shorter than expected for the given number of levels.
-		// Handle this case appropriately, either by returning an error or by handling it in some other way.
-		panic(fmt.Sprintf("key length is too short: expected at least %d bytes, got %d bytes", requiredLen, len(k)))
+		padding := make([]byte, requiredLen-len(k))
+		k = append(k, padding...)
 	}
 
 	path := make([]bool, numLevels)

tree/arbo/tree.go

@@ -157,13 +157,7 @@ func NewTreeWithTx(wTx db.WriteTx, cfg Config) (*Tree, error) {
 	_, err = wTx.Get(dbKeyRoot)
 	if err == db.ErrKeyNotFound {
 		// store new root 0 (empty)
-		if err = wTx.Set(dbKeyRoot, t.emptyHash); err != nil {
-			return nil, err
-		}
-		if err = t.setNLeafs(wTx, 0); err != nil {
-			return nil, err
-		}
-		return &t, nil
+		return &t, t.setToEmptyTree(wTx)
 	} else if err != nil {
 		return nil, fmt.Errorf("unknown database error: %w", err)
 	}
@@ -286,9 +280,11 @@ func (t *Tree) add(wTx db.WriteTx, root []byte, fromLvl int, k, v []byte) ([]byt
 
 	// go up to the root
 	if len(siblings) == 0 {
-		// return the leafKey as root
+		// if there are no siblings, means that the tree is empty.
+		// We consider the leafKey as root
 		return leafKey, nil
 	}
+
 	root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, fromLvl)
 	if err != nil {
 		return nil, err
@@ -425,6 +421,14 @@ func (t *Tree) DeleteWithTx(wTx db.WriteTx, k []byte) error {
 	return t.deleteWithTx(wTx, k)
 }
 
+// setToEmptyTree sets the tree to an empty tree.
+func (t *Tree) setToEmptyTree(wTx db.WriteTx) error {
+	if err := wTx.Set(dbKeyRoot, t.emptyHash); err != nil {
+		return err
+	}
+	return t.setNLeafs(wTx, 0)
+}
+
 // deleteWithTx deletes a key-value pair from the tree.
 func (t *Tree) deleteWithTx(wTx db.WriteTx, k []byte) error {
 	root, err := t.RootWithTx(wTx)
@@ -445,12 +449,17 @@ func (t *Tree) deleteWithTx(wTx db.WriteTx, k []byte) error {
 		return err
 	}
 
-	// if the neighbor is not empty, set the leaf key to the empty hash
 	if len(siblings) == 0 {
-		return nil
+		// if no siblings, means the tree is now empty,
+		// just delete the leaf key and set the root to empty
+		if err := wTx.Delete(leafKey); err != nil {
+			return fmt.Errorf("error deleting leaf key %x: %w", leafKey, err)
+		}
+		return t.setToEmptyTree(wTx)
 	}
 
 	var neighbourKeys, neighbourValues [][]byte
+	// if the neighbour is not empty, we set it to empty (instead of remove)
 	if !bytes.Equal(siblings[len(siblings)-1], t.emptyHash) {
 		neighbourKeys, neighbourValues, err = t.getLeavesFromSubPath(wTx, siblings[len(siblings)-1])
 		if err != nil {
@@ -470,10 +479,6 @@ func (t *Tree) deleteWithTx(wTx db.WriteTx, k []byte) error {
 	}
 
 	// Navigate back up the tree, updating the intermediate nodes.
-	if len(siblings) == 0 {
-		// The tree is empty after deletion.
-		return t.setRoot(wTx, t.emptyHash)
-	}
 	newRoot, err := t.up(wTx, t.emptyHash, siblings, path, len(siblings)-1, 0)
 	if err != nil {
 		return err
@@ -489,7 +494,7 @@ func (t *Tree) deleteWithTx(wTx db.WriteTx, k []byte) error {
 		return fmt.Errorf("error deleting orphan intermediate nodes: %v", err)
 	}
 
-	// If there are neighbours deleted, add them back to the tree.
+	// Delete the neighbour's childs and add them back to the tree in the right place.
 	for i, k := range neighbourKeys {
 		if err := t.deleteWithTx(wTx, k); err != nil {
 			return fmt.Errorf("error deleting neighbour %d: %w", i, err)

tree/arbo/tree_test.go

@@ -947,7 +947,7 @@ func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
 func TestDiskSizeBench(t *testing.T) {
 	c := qt.New(t)
 
-	nLeafs := 2000
+	nLeafs := 500
 	printTestContext("TestDiskSizeBench: ", nLeafs, "Blake2b", "pebble")
 
 	// prepare inputs
@@ -1204,3 +1204,38 @@ func contains(arr [][]byte, item []byte) bool {
 	}
 	return false
 }
+
+func TestTreeWithSingleLeaf(t *testing.T) {
+	c := qt.New(t)
+
+	database := metadb.NewTest(t)
+	tree, err := NewTree(Config{Database: database, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
+	c.Assert(err, qt.IsNil)
+
+	// check empty root
+	root, err := tree.Root()
+	c.Assert(err, qt.IsNil)
+	c.Assert(hex.EncodeToString(root), qt.DeepEquals, "0000000000000000000000000000000000000000000000000000000000000000")
+
+	// add one entry
+	err = tree.Add([]byte{0x01}, []byte{0x01})
+	c.Assert(err, qt.IsNil)
+
+	root, err = tree.Root()
+	c.Assert(err, qt.IsNil)
+	c.Assert(root, qt.HasLen, 32)
+
+	leafKey, _, err := newLeafValue(HashFunctionPoseidon, []byte{0x01}, []byte{0x01})
+	c.Assert(err, qt.IsNil)
+
+	// check leaf key is actually the current root
+	c.Assert(bytes.Equal(root, leafKey), qt.IsTrue)
+
+	// delete the only entry and check the tree is empty again
+	err = tree.Delete([]byte{0x01})
+	c.Assert(err, qt.IsNil)
+	root, err = tree.Root()
+	c.Assert(err, qt.IsNil)
+	c.Assert(hex.EncodeToString(root), qt.DeepEquals, "0000000000000000000000000000000000000000000000000000000000000000")
+}

vochain/state/sik_test.go

@@ -159,13 +159,14 @@ func TestAssignSIKToElectionAndPurge(t *testing.T) {
 	c.Assert(err, qt.IsNil)
 	_, err = s.NoState(true).Get(key)
 	c.Assert(err, qt.IsNil)
+
 	// purge siks and check
 	c.Assert(s.PurgeSIKsByElection(pid), qt.IsNil)
 	c.Assert(err, qt.IsNil)
 	_, err = s.NoState(true).Get(key)
 	c.Assert(err, qt.IsNotNil)
-	_, err = s.SIKFromAddress(testAccount.Address())
-	c.Assert(err, qt.IsNotNil)
+	sik, err = s.SIKFromAddress(testAccount.Address())
+	c.Assert(err, qt.IsNotNil, qt.Commentf("SIK should be deleted"))
 }
 
 func Test_heightEncoding(t *testing.T) {