Skip to content

Commit 3623225

Browse files
committed
Add the ability to decode a single path
This is more ergonomic than creating a struct for a single value and the performance is better as well due to less reflection: BenchmarkDecodeCountryCodeWithStruct-8 1347441 882.4 ns/op 1 B/op 0 allocs/op BenchmarkDecodePathCountryCode-8 2708011 445.1 ns/op 1 B/op 0 allocs/op
1 parent 601a682 commit 3623225

File tree

3 files changed

+167
-1
lines changed

3 files changed

+167
-1
lines changed

decoder.go

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,82 @@ func (d *decoder) decodeToDeserializer(
8989
return d.decodeFromTypeToDeserializer(typeNum, size, newOffset, dser, depth+1)
9090
}
9191

92+
func (d *decoder) decodePath(
93+
offset uint,
94+
path []any,
95+
result reflect.Value,
96+
) error {
97+
PATH:
98+
for i, v := range path {
99+
var (
100+
typeNum dataType
101+
size uint
102+
err error
103+
)
104+
typeNum, size, offset, err = d.decodeCtrlData(offset)
105+
if err != nil {
106+
return err
107+
}
108+
109+
if typeNum == _Pointer {
110+
pointer, _, err := d.decodePointer(size, offset)
111+
if err != nil {
112+
return err
113+
}
114+
115+
typeNum, size, offset, err = d.decodeCtrlData(pointer)
116+
if err != nil {
117+
return err
118+
}
119+
}
120+
121+
switch v := v.(type) {
122+
case string:
123+
// We are expecting a map
124+
if typeNum != _Map {
125+
// XXX - use type names in errors.
126+
return fmt.Errorf("expected a map for %s but found %d", v, typeNum)
127+
}
128+
for i := uint(0); i < size; i++ {
129+
var key []byte
130+
key, offset, err = d.decodeKey(offset)
131+
if err != nil {
132+
return err
133+
}
134+
if string(key) == v {
135+
continue PATH
136+
}
137+
offset, err = d.nextValueOffset(offset, 1)
138+
if err != nil {
139+
return err
140+
}
141+
}
142+
// Not found. Maybe return a boolean?
143+
return nil
144+
case int:
145+
// We are expecting an array
146+
if typeNum != _Slice {
147+
// XXX - use type names in errors.
148+
return fmt.Errorf("expected a slice for %d but found %d", v, typeNum)
149+
}
150+
if size < uint(v) {
151+
// Slice is smaller than index, not found
152+
return nil
153+
}
154+
// TODO: support negative indexes? Seems useful for subdivisions in
155+
// particular.
156+
offset, err = d.nextValueOffset(offset, uint(v))
157+
if err != nil {
158+
return err
159+
}
160+
default:
161+
return fmt.Errorf("unexpected type for %d value in path, %v: %T", i, v, v)
162+
}
163+
}
164+
_, err := d.decode(offset, result, len(path))
165+
return err
166+
}
167+
92168
func (d *decoder) decodeCtrlData(offset uint) (dataType, uint, uint, error) {
93169
newOffset := offset + 1
94170
if offset >= uint(len(d.buffer)) {

reader_test.go

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,30 @@ func TestDecoder(t *testing.T) {
316316
require.NoError(t, reader.Close())
317317
}
318318

319+
func TestDecodePath(t *testing.T) {
320+
reader, err := Open(testFile("MaxMind-DB-test-decoder.mmdb"))
321+
require.NoError(t, err)
322+
323+
result := reader.Lookup(netip.MustParseAddr("::1.1.1.0"))
324+
require.NoError(t, result.Err())
325+
326+
var u16 uint16
327+
328+
require.NoError(t, result.DecodePath(&u16, "uint16"))
329+
330+
assert.Equal(t, uint16(100), u16)
331+
332+
var u uint
333+
require.NoError(t, result.DecodePath(&u, "array", 0))
334+
assert.Equal(t, uint(1), u)
335+
336+
require.NoError(t, result.DecodePath(&u, "array", 2))
337+
assert.Equal(t, uint(3), u)
338+
339+
require.NoError(t, result.DecodePath(&u, "map", "mapX", "arrayX", 1))
340+
assert.Equal(t, uint(8), u)
341+
}
342+
319343
type TestInterface interface {
320344
method() bool
321345
}
@@ -902,7 +926,7 @@ func BenchmarkCityLookupNetwork(b *testing.B) {
902926
require.NoError(b, db.Close(), "error on close")
903927
}
904928

905-
func BenchmarkCountryCode(b *testing.B) {
929+
func BenchmarkDecodeCountryCodeWithStruct(b *testing.B) {
906930
db, err := Open("GeoLite2-City.mmdb")
907931
require.NoError(b, err)
908932

@@ -927,6 +951,27 @@ func BenchmarkCountryCode(b *testing.B) {
927951
require.NoError(b, db.Close(), "error on close")
928952
}
929953

954+
func BenchmarkDecodePathCountryCode(b *testing.B) {
955+
db, err := Open("GeoLite2-City.mmdb")
956+
require.NoError(b, err)
957+
958+
path := []any{"country", "iso_code"}
959+
960+
//nolint:gosec // this is a test
961+
r := rand.New(rand.NewSource(0))
962+
var result string
963+
964+
s := make(net.IP, 4)
965+
for i := 0; i < b.N; i++ {
966+
ip := randomIPv4Address(r, s)
967+
err = db.Lookup(ip).DecodePath(&result, path...)
968+
if err != nil {
969+
b.Error(err)
970+
}
971+
}
972+
require.NoError(b, db.Close(), "error on close")
973+
}
974+
930975
func randomIPv4Address(r *rand.Rand, ip []byte) netip.Addr {
931976
num := r.Uint32()
932977
ip[0] = byte(num >> 24)

result.go

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,51 @@ func (r Result) Decode(v any) error {
4646
return err
4747
}
4848

49+
// DecodePath unmarshals a value from data section into v, following the
50+
// specified path.
51+
//
52+
// The v parameter should be a pointer to the value where the decoded data
53+
// will be stored. If v is nil or not a pointer, an error is returned. If the
54+
// data in the database record cannot be stored in v because of type
55+
// differences, an UnmarshalTypeError is returned.
56+
//
57+
// The path is a variadic list of keys (strings) and/or indices (ints) that
58+
// describe the nested structure to traverse in the data to reach the desired
59+
// value.
60+
//
61+
// For maps, string path elements are used as keys.
62+
// For arrays, int path elements are used as indices.
63+
//
64+
// If the path is empty, the entire data structure is decoded into v.
65+
//
66+
// Returns an error if:
67+
// - the path is invalid
68+
// - the data cannot be decoded into the type of v
69+
// - v is not a pointer
70+
// - the database record cannot be stored in v because of the type
71+
// - the Result does not contain valid data
72+
//
73+
// Example usage:
74+
//
75+
// var city string
76+
// err := result.DecodePath(&city, "location", "city", "names", "en")
77+
//
78+
// var geonameID int
79+
// err := result.DecodePath(&geonameID, "subdivisions", 0, "geoname_id")
80+
func (r Result) DecodePath(v any, path ...any) error {
81+
if r.err != nil {
82+
return r.err
83+
}
84+
if r.offset == notFound {
85+
return nil
86+
}
87+
rv := reflect.ValueOf(v)
88+
if rv.Kind() != reflect.Ptr || rv.IsNil() {
89+
return errors.New("result param must be a pointer")
90+
}
91+
return r.decoder.decodePath(r.offset, path, rv)
92+
}
93+
4994
// Err provides a way to check whether there was an error during the lookup
5095
// without clling Result.Decode. If there was an error, it will also be
5196
// returned from Result.Decode.

0 commit comments

Comments
 (0)