diff --git a/db/trie/branchroottrie.go b/db/trie/branchroottrie.go index e0f39514ee..1173acd12d 100644 --- a/db/trie/branchroottrie.go +++ b/db/trie/branchroottrie.go @@ -74,6 +74,10 @@ func (tr *branchRootTrie) SetRootHash(rootHash []byte) error { return nil } +func (tr *branchRootTrie) IsEmpty() bool { + return tr.isEmptyRootHash(tr.rootHash) +} + func (tr *branchRootTrie) Get(key []byte) ([]byte, error) { trieMtc.WithLabelValues("root", "Get").Inc() kt, err := tr.checkKeyType(key) diff --git a/db/trie/trie.go b/db/trie/trie.go index af9adf43bf..63b1512380 100644 --- a/db/trie/trie.go +++ b/db/trie/trie.go @@ -58,6 +58,8 @@ type Trie interface { RootHash() []byte // SetRootHash sets a new root to trie SetRootHash([]byte) error + // IsEmpty returns true is this is an empty trie + IsEmpty() bool // DB returns the KVStore storing the node data DB() KVStore // deleteNodeFromDB deletes the data of node from db diff --git a/state/factory/twolayertrie.go b/state/factory/twolayertrie.go new file mode 100644 index 0000000000..5142b69ba5 --- /dev/null +++ b/state/factory/twolayertrie.go @@ -0,0 +1,100 @@ +// Copyright (c) 2019 IoTeX Foundation +// This is an alpha (internal) release and is not suitable for production. This source code is provided 'as is' and no +// warranties are given as to title or non-infringement, merchantability or fitness for purpose and, to the extent +// permitted by law, all liability for your use of the code is disclaimed. This source code is governed by Apache +// License 2.0 that can be found in the LICENSE file. + +package factory + +import ( + "context" + + "github.com/iotexproject/iotex-core/db/trie" + "github.com/pkg/errors" +) + +// TwoLayerTrie is a trie data structure with two layers +type TwoLayerTrie struct { + layerOne trie.Trie +} + +// Start starts the layer one trie +func (tlt *TwoLayerTrie) Start(ctx context.Context) error { + return tlt.layerOne.Start(ctx) +} + +// Stop stops the layer one trie +func (tlt *TwoLayerTrie) Stop(ctx context.Context) error { + return tlt.layerOne.Stop(ctx) +} + +// RootHash returns the layer one trie root +func (tlt *TwoLayerTrie) RootHash() []byte { + return tlt.layerOne.RootHash() +} + +func (tlt *TwoLayerTrie) layerTwoTrie(key []byte) (trie.Trie, error) { + value, err := tlt.layerOne.Get(key) + switch errors.Cause(err) { + case trie.ErrNotExist: + return trie.NewTrie(trie.KVStoreOption(tlt.layerOne.DB()), trie.KeyLengthOption(len(key))) + case nil: + return trie.NewTrie(trie.KVStoreOption(tlt.layerOne.DB()), trie.RootHashOption(value), trie.KeyLengthOption(len(key))) + default: + return nil, err + } +} + +// Get returns the layer two value +func (tlt *TwoLayerTrie) Get(layerOneKey []byte, layerTwoKey []byte) ([]byte, error) { + layerTwo, err := tlt.layerTwoTrie(layerOneKey) + if err != nil { + return nil, err + } + if err := layerTwo.Start(context.Background()); err != nil { + return nil, err + } + defer layerTwo.Stop(context.Background()) + + return layerTwo.Get(layerTwoKey) +} + +// Upsert upserts an item +func (tlt *TwoLayerTrie) Upsert(layerOneKey []byte, layerTwoKey []byte, value []byte) error { + layerTwo, err := tlt.layerTwoTrie(layerOneKey) + if err != nil { + return err + } + if err := layerTwo.Start(context.Background()); err != nil { + return err + } + defer layerTwo.Stop(context.Background()) + + if err := layerTwo.Upsert(layerTwoKey, value); err != nil { + return err + } + + return tlt.layerOne.Upsert(layerOneKey, layerTwo.RootHash()) +} + +// Delete deletes an item +func (tlt *TwoLayerTrie) Delete(layerOneKey []byte, layerTwoKey []byte) error { + layerTwo, err := tlt.layerTwoTrie(layerOneKey) + if err != nil { + return err + } + if err := layerTwo.Start(context.Background()); err != nil { + return err + } + defer layerTwo.Stop(context.Background()) + + if err := layerTwo.Delete(layerTwoKey); err != nil { + return err + } + + if layerTwo.IsEmpty() { + return tlt.layerOne.Delete(layerOneKey) + } + + return tlt.layerOne.Upsert(layerOneKey, layerTwo.RootHash()) +} diff --git a/state/factory/twolayertrie_test.go b/state/factory/twolayertrie_test.go new file mode 100644 index 0000000000..3327d3621f --- /dev/null +++ b/state/factory/twolayertrie_test.go @@ -0,0 +1,9 @@ +// Copyright (c) 2019 IoTeX Foundation +// This is an alpha (internal) release and is not suitable for production. This source code is provided 'as is' and no +// warranties are given as to title or non-infringement, merchantability or fitness for purpose and, to the extent +// permitted by law, all liability for your use of the code is disclaimed. This source code is governed by Apache +// License 2.0 that can be found in the LICENSE file. + +package factory + +// TODO: add unit test for two layer trie diff --git a/test/mock/mock_factory/mock_workingset.go b/test/mock/mock_factory/mock_workingset.go index a53eae9059..4cb96824e4 100644 --- a/test/mock/mock_factory/mock_workingset.go +++ b/test/mock/mock_factory/mock_workingset.go @@ -212,10 +212,10 @@ func (mr *MockWorkingSetMockRecorder) Commit() *gomock.Call { } // RootHash mocks base method -func (m *MockWorkingSet) RootHash() (hash.Hash256, error) { +func (m *MockWorkingSet) RootHash() ([]byte, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "RootHash") - ret0, _ := ret[0].(hash.Hash256) + ret0, _ := ret[0].([]byte) ret1, _ := ret[1].(error) return ret0, ret1 } diff --git a/test/mock/mock_trie/mock_trie.go b/test/mock/mock_trie/mock_trie.go index 985af611fb..4d7b3d759d 100644 --- a/test/mock/mock_trie/mock_trie.go +++ b/test/mock/mock_trie/mock_trie.go @@ -133,6 +133,20 @@ func (mr *MockTrieMockRecorder) SetRootHash(arg0 interface{}) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetRootHash", reflect.TypeOf((*MockTrie)(nil).SetRootHash), arg0) } +// IsEmpty mocks base method +func (m *MockTrie) IsEmpty() bool { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "IsEmpty") + ret0, _ := ret[0].(bool) + return ret0 +} + +// IsEmpty indicates an expected call of IsEmpty +func (mr *MockTrieMockRecorder) IsEmpty() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsEmpty", reflect.TypeOf((*MockTrie)(nil).IsEmpty)) +} + // DB mocks base method func (m *MockTrie) DB() trie.KVStore { m.ctrl.T.Helper()