diff --git a/art_node.go b/art_node.go index 2cf3eb4..18305e1 100644 --- a/art_node.go +++ b/art_node.go @@ -45,19 +45,27 @@ type ArtNode struct { size uint8 // Leaf Node Attributes - key []byte - keySize uint64 - value interface{} - nodeType uint8 + key []byte + keySize uint64 + completeKey []byte + value interface{} + nodeType uint8 } func NewLeafNode(key []byte, value interface{}) *ArtNode { + return newLeafNode(key, key, value) +} + +func newLeafNode(key, completeKey []byte, value interface{}) *ArtNode { newKey := make([]byte, len(key)) copy(newKey, key) + originalKey := make([]byte, len(completeKey)) + copy(originalKey, completeKey) l := &ArtNode{ - key: newKey, - value: value, - nodeType: LEAF, + key: newKey, + completeKey: originalKey, + value: value, + nodeType: LEAF, } return l @@ -647,6 +655,15 @@ func (n *ArtNode) copyMeta(other *ArtNode) { n.prefixLen = other.prefixLen } +// Key() returns the key used when the node was inserted +func (n *ArtNode) Key() []byte { + if n.nodeType != LEAF { + return nil + } + + return n.completeKey +} + // Returns the value of the given node, or nil if it is not a leaf. func (n *ArtNode) Value() interface{} { if n.nodeType != LEAF { diff --git a/art_node_test.go b/art_node_test.go index 11a2bbb..4cffb64 100644 --- a/art_node_test.go +++ b/art_node_test.go @@ -215,7 +215,7 @@ func TestShrink(t *testing.T) { } func TestNewLeafNode(t *testing.T) { - key := []byte{'a', 'r', 't' } + key := []byte{'a', 'r', 't'} value := "tree" l := NewLeafNode(key, value) @@ -227,6 +227,10 @@ func TestNewLeafNode(t *testing.T) { t.Errorf("Expected key value to match the one supplied") } + if bytes.Compare(l.completeKey, key) != 0 { + t.Errorf("Expected completeKey value to match the one supplied") + } + if l.value != value { t.Errorf("Expected initial value to match the one supplied") } @@ -234,4 +238,4 @@ func TestNewLeafNode(t *testing.T) { if l.nodeType != LEAF { t.Errorf("Expected Leaf node to be of LEAF type") } -} \ No newline at end of file +} diff --git a/art_tree.go b/art_tree.go index b278978..2b5117a 100644 --- a/art_tree.go +++ b/art_tree.go @@ -58,8 +58,9 @@ func (t *ArtTree) searchHelper(current *ArtNode, key []byte, depth int) interfac // Inserts the passed in value that is indexed by the passed in key into the ArtTree. func (t *ArtTree) Insert(key []byte, value interface{}) { + var completeKey = key; key = ensureNullTerminatedKey(key) - t.insertHelper(t.root, &t.root, key, value, 0) + t.insertHelper(t.root, &t.root, key, completeKey, value, 0) } // Recursive helper function that traverses the tree until an insertion point is found. @@ -77,12 +78,12 @@ func (t *ArtTree) Insert(key []byte, value interface{}) { // // If there is no child at the specified key at the current depth of traversal, a new leaf node // is created and inserted at this position. -func (t *ArtTree) insertHelper(current *ArtNode, currentRef **ArtNode, key []byte, value interface{}, depth int) { +func (t *ArtTree) insertHelper(current *ArtNode, currentRef **ArtNode, key, completeKey []byte, value interface{}, depth int) { // @spec: Usually, the leaf can // simply be inserted into an existing inner node, after growing // it if necessary. if current == nil { - *currentRef = NewLeafNode(key, value) + *currentRef = newLeafNode(key, completeKey, value) t.size += 1 return } @@ -100,7 +101,7 @@ func (t *ArtTree) insertHelper(current *ArtNode, currentRef **ArtNode, key []byt // Create a new Inner Node to contain the new Leaf and the current node. newNode4 := NewNode4() - newLeafNode := NewLeafNode(key, value) + newLeafNode := newLeafNode(key, completeKey, value) // Determine the longest common prefix between our current node and the key limit := current.LongestCommonPrefix(newLeafNode, depth) @@ -150,7 +151,7 @@ func (t *ArtTree) insertHelper(current *ArtNode, currentRef **ArtNode, key []byt } // Attach the desired insertion key - newLeafNode := NewLeafNode(key, value) + newLeafNode := newLeafNode(key, completeKey, value) newNode4.AddChild(key[depth+mismatch], newLeafNode) t.size += 1 @@ -167,11 +168,11 @@ func (t *ArtTree) insertHelper(current *ArtNode, currentRef **ArtNode, key []byt if *next != nil { // Recurse, and keep looking for an insertion point - t.insertHelper(*next, next, key, value, depth+1) + t.insertHelper(*next, next, key, completeKey, value, depth+1) } else { // Otherwise, Add the child at the current position. - current.AddChild(key[depth], NewLeafNode(key, value)) + current.AddChild(key[depth], newLeafNode(key, completeKey, value)) t.size += 1 } }