Delete code with account, add code test (#171)

* Delete code with account, add code test

* Fix test

Co-authored-by: Alexey Sharp <alexeysharp@Alexeys-iMac.local>

aggregator/aggregator.go

@@ -948,7 +948,7 @@ func (r *Reader) ReadAccountStorage(addr []byte, loc []byte) (*uint256.Int, erro
 	return nil, nil
 }
 
-func (r *Reader) ReadAccountCode(addr []byte, incarnation uint64) ([]byte, error) {
+func (r *Reader) ReadAccountCode(addr []byte) ([]byte, error) {
 	// Look in the summary table first
 	v, err := r.tx.GetOne(kv.StateCode, addr)
 	if err != nil {
@@ -963,10 +963,6 @@ func (r *Reader) ReadAccountCode(addr []byte, incarnation uint64) ([]byte, error
 	return val, nil
 }
 
-func (r *Reader) ReadAccountIncarnation(addr []byte) uint64 {
-	return r.blockNum
-}
-
 type Writer struct {
 	a        *Aggregator
 	tx       kv.RwTx
@@ -1138,7 +1134,7 @@ func (ch *CursorHeap) Pop() interface{} {
 	return x
 }
 
-func (w *Writer) DeleteAccount(addr []byte) error {
+func (w *Writer) deleteAccount(addr []byte) error {
 	prevV, err := w.tx.GetOne(kv.StateAccounts, addr)
 	if err != nil {
 		return err
@@ -1163,14 +1159,53 @@ func (w *Writer) DeleteAccount(addr []byte) error {
 	if err = w.a.accountChanges.delete(addr, original); err != nil {
 		return err
 	}
+	return nil
+}
+
+func (w *Writer) deleteCode(addr []byte) error {
+	prevV, err := w.tx.GetOne(kv.StateCode, addr)
+	if err != nil {
+		return err
+	}
+	var prevNum uint32
+	var original []byte
+	if prevV == nil {
+		original = w.a.readCode(w.blockNum, addr)
+	} else {
+		prevNum = binary.BigEndian.Uint32(prevV[:4])
+	}
+	v := make([]byte, 4)
+	binary.BigEndian.PutUint32(v[:4], prevNum+1)
+	if err = w.tx.Put(kv.StateCode, addr, v); err != nil {
+		return err
+	}
+	if prevV == nil && original == nil {
+		// Nothing to do
+		return nil
+	} else if original == nil {
+		original = prevV[4:]
+	}
+	if err = w.a.codeChanges.delete(addr, original); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (w *Writer) DeleteAccount(addr []byte) error {
+	if err := w.deleteAccount(addr); err != nil {
+		return err
+	}
+	if err := w.deleteCode(addr); err != nil {
+		return err
+	}
 	// Find all storage items for this address
 	var cp CursorHeap
 	heap.Init(&cp)
-	var c kv.Cursor
-	if c, err = w.tx.Cursor(kv.StateStorage); err != nil {
+	c, err := w.tx.Cursor(kv.StateStorage)
+	if err != nil {
 		return err
 	}
-	var k []byte
+	var k, v []byte
 	if k, v, err = c.Seek(addr); err != nil {
 		return err
 	}
@@ -1179,6 +1214,9 @@ func (w *Writer) DeleteAccount(addr []byte) error {
 	}
 	w.a.byEndBlock.Ascend(func(i btree.Item) bool {
 		item := i.(*byEndBlockItem)
+		if item.storageIdx.Empty() {
+			return true
+		}
 		offset := item.storageIdx.Lookup(addr)
 		g := item.storageD.MakeGetter() // TODO Cache in the reader
 		g.Reset(offset)
@@ -1230,11 +1268,12 @@ func (w *Writer) DeleteAccount(addr []byte) error {
 				}
 			}
 		}
+		var prevV []byte
 		prevV, err = w.tx.GetOne(kv.StateStorage, lastKey)
 		if err != nil {
 			return err
 		}
-		prevNum = 0
+		var prevNum uint32
 		if prevV != nil {
 			prevNum = binary.BigEndian.Uint32(prevV[:4])
 		}

aggregator/aggregator_test.go

@@ -279,3 +279,86 @@ func TestRecreateAccountWithStorage(t *testing.T) {
 		tx.Rollback()
 	}
 }
+
+func TestChangeCode(t *testing.T) {
+	tmpDir := t.TempDir()
+	db := memdb.New()
+	defer db.Close()
+	a, err := NewAggregator(tmpDir, 16, 4)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer a.Close()
+	accountKey := int160(1)
+	var account1 = int256(1)
+	var code1 = []byte("This is the code number 1")
+	//var code2 = []byte("This is the code number 2")
+	var rwTx kv.RwTx
+	defer func() {
+		rwTx.Rollback()
+	}()
+	var tx kv.Tx
+	defer func() {
+		if tx != nil {
+			tx.Rollback()
+		}
+	}()
+	for blockNum := uint64(0); blockNum < 100; blockNum++ {
+		if rwTx, err = db.BeginRw(context.Background()); err != nil {
+			t.Fatal(err)
+		}
+		var w *Writer
+		if w, err = a.MakeStateWriter(rwTx, blockNum); err != nil {
+			t.Fatal(err)
+		}
+		switch blockNum {
+		case 1:
+			if err = w.UpdateAccountData(accountKey, account1); err != nil {
+				t.Fatal(err)
+			}
+			if err = w.UpdateAccountCode(accountKey, code1); err != nil {
+				t.Fatal(err)
+			}
+		case 25:
+			if err = w.DeleteAccount(accountKey); err != nil {
+				t.Fatal(err)
+			}
+		}
+		if err = w.Finish(); err != nil {
+			t.Fatal(err)
+		}
+		if err = rwTx.Commit(); err != nil {
+			t.Fatal(err)
+		}
+		if tx, err = db.BeginRo(context.Background()); err != nil {
+			t.Fatal(err)
+		}
+		r := a.MakeStateReader(tx, blockNum+1)
+		switch blockNum {
+		case 22:
+			var acc []byte
+			if acc, err = r.ReadAccountData(accountKey); err != nil {
+				t.Fatal(err)
+			}
+			if !bytes.Equal(account1, acc) {
+				t.Errorf("wrong account after block %d, expected %x, got %x", blockNum, account1, acc)
+			}
+			var code []byte
+			if code, err = r.ReadAccountCode(accountKey); err != nil {
+				t.Fatal(err)
+			}
+			if !bytes.Equal(code1, code) {
+				t.Errorf("wrong code after block %d, expected %x, got %x", blockNum, code1, code)
+			}
+		case 47:
+			var code []byte
+			if code, err = r.ReadAccountCode(accountKey); err != nil {
+				t.Fatal(err)
+			}
+			if code != nil {
+				t.Errorf("wrong code after block %d, expected nil, got %x", blockNum, code)
+			}
+		}
+		tx.Rollback()
+	}
+}