Allow `Tree` to have default insert func
UmanShahzad opened this issue · 1 comments
UmanShahzad commented
Right now the Tree.Insert
function defaults to the inserter.ReplaceWith
strategy. You have to use Tree.InsertFunc
to use any other strategy.
It would be nice if we could choose the default for Tree.Insert
. Something close to what this diff achieves:
diff --git a/tree.go b/tree.go
index 4e436a8..ce0fca9 100644
--- a/tree.go
+++ b/tree.go
@@ -72,6 +72,11 @@ type Options struct {
// implementations that do not correctly handle metadata pointers. Its
// use should primarily be limited to existing database types.
DisableMetadataPointers bool
+
+ // Inserter is the insert function used when calling `Insert`. It defaults
+ // to `inserter.ReplaceWith`, which replaces any conflicting old value
+ // entirely with the new.
+ Inserter func(value mmdbtype.DataType) (mmdbtype.DataType, error)
}
// Tree represents an MaxMind DB search tree.
@@ -88,6 +93,7 @@ type Tree struct {
treeDepth int
// This is set when the tree is finalized
nodeCount int
+ inserterFunc func(value mmdbtype.DataType) (mmdbtype.DataType, error)
}
// New creates a new Tree.
@@ -101,6 +107,7 @@ func New(opts Options) (*Tree, error) {
ipVersion: 6,
recordSize: 28,
root: &node{},
+ inserterFunc: inserter.ReplaceWith,
}
if opts.BuildEpoch != 0 {
@@ -123,6 +130,10 @@ func New(opts Options) (*Tree, error) {
tree.recordSize = opts.RecordSize
}
+ if opts.Inserter != nil {
+ tree.inserterFunc = opts.Inserter
+ }
+
switch tree.ipVersion {
case 6:
tree.treeDepth = 128
@@ -214,7 +225,7 @@ func Load(path string, opts Options) (*Tree, error) {
//
// This is not safe to call from multiple threads.
func (t *Tree) Insert(network *net.IPNet, value mmdbtype.DataType) error {
- return t.InsertFunc(network, inserter.ReplaceWith(value))
+ return t.InsertFunc(network, t.inserterFunc(value))
}
// InsertFunc will insert the output of the function passed to it. The argument
oschwald commented
If you made a PR, I think it would be accepted. The documentation for Insert
would probably need to be updated slightly as well.