Share traverseTree

maxmind · Nov 10, 2023 · 7c6a5e1 · 7c6a5e1
1 parent ae8ca1c
commit 7c6a5e1
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 54 deletions.
diff --git a/src/main/java/com/maxmind/db/Reader.java b/src/main/java/com/maxmind/db/Reader.java
@@ -169,21 +169,16 @@ int getIpv4Start() {
      */
     public <T> DatabaseRecord<T> getRecord(InetAddress ipAddress, Class<T> cls)
         throws IOException {
-        ByteBuffer buffer = this.getBufferHolder().get();
 
         byte[] rawAddress = ipAddress.getAddress();
 
-        int bitLength = rawAddress.length * 8;
-        int record = this.startNode(bitLength);
-        int nodeCount = this.metadata.getNodeCount();
+        int[] traverseResult = traverseTree(rawAddress, rawAddress.length * 8);
 
-        int pl = 0;
-        for (; pl < bitLength && record < nodeCount; pl++) {
-            int b = 0xFF & rawAddress[pl / 8];
-            int bit = 1 & (b >> 7 - (pl % 8));
-            record = this.readNode(buffer, record, bit);
-        }
+        int pl = traverseResult[1];
+        int record = traverseResult[0];
 
+        int nodeCount = this.metadata.getNodeCount();
+        ByteBuffer buffer = this.getBufferHolder().get();
         T dataRecord = null;
         if (record > nodeCount) {
             // record is a data pointer
@@ -317,7 +312,7 @@ public <T> Networks<T> networksWithin(
             prefixLength += 96;
         }
 
-        int[] traverseResult = this.traverseTree(ipBytes, 0, prefixLength);
+        int[] traverseResult = this.traverseTree(ipBytes, prefixLength);
         int node = traverseResult[0];
         int prefix = traverseResult[1];
 
@@ -335,22 +330,24 @@ public <T> Networks<T> networksWithin(
      * @param bitCount The prefix.
      * @return int[]
      */
-    private int[] traverseTree(byte[] ip, int node, int bitCount)
+    private int[] traverseTree(byte[] ip, int bitCount)
         throws ClosedDatabaseException, InvalidDatabaseException {
-        int nodeCount = this.metadata.getNodeCount();
-        int i = 0;
-
         ByteBuffer buffer = this.getBufferHolder().get();
+        int bitLength = ip.length * 8;
+        int record = this.startNode(bitLength);
+        int nodeCount = this.metadata.getNodeCount();
 
-        for (; i < bitCount && node < nodeCount; i++) {
-            int bit = 1 & (ip[i >> 3] >> (7 - (i % 8)));
+        int i = 0;
+        for (; i < bitCount && record < nodeCount; i++) {
+            int b = 0xFF & ip[i / 8];
+            int bit = 1 & (b >> 7 - (i % 8));
 
             // bit:0 -> left record.
             // bit:1 -> right record.
-            node = this.readNode(buffer, node, bit);
+            record = this.readNode(buffer, record, bit);
         }
 
-        return new int[]{node, i};
+        return new int[]{record, i};
     }
 
     int readNode(ByteBuffer buffer, int nodeNumber, int index)

diff --git a/src/test/java/com/maxmind/db/ReaderTest.java b/src/test/java/com/maxmind/db/ReaderTest.java
@@ -290,41 +290,41 @@ public networkTest(String network,  int prefix,String database, String[] expecte
                 "2002:101:120:0:0:0:0:0/48",
             }
         ),
-        // new networkTest(
-        //     "::",
-        //     0,
-        //     "mixed",
-        //     new String[]{
-        //         "1.1.1.1/32",
-        //         "1.1.1.2/31",
-        //         "1.1.1.4/30",
-        //         "1.1.1.8/29",
-        //         "1.1.1.16/28",
-        //         "1.1.1.32/32",
-        //         "0:0:0:0:0:1:ffff:ffff/128",
-        //         "0:0:0:0:0:2:0:0/122",
-        //         "0:0:0:0:0:2:0:40/124",
-        //         "0:0:0:0:0:2:0:50/125",
-        //         "0:0:0:0:0:2:0:58/127",
-        //     },
-        //     true
-        // ),
-        // new networkTest(
-        //     "1.1.1.16",
-        //     28,
-        //     "mixed",
-        //     new String[]{
-        //         "1.1.1.16/28"
-        //     }
-        // ),
-        // new networkTest(
-        //     "1.1.1.4",
-        //     30,
-        //     "ipv4",
-        //     new String[]{
-        //         "1.1.1.4/30"
-        //     }
-        // )
+        new networkTest(
+            "::",
+            0,
+            "mixed",
+            new String[]{
+                "1.1.1.1/32",
+                "1.1.1.2/31",
+                "1.1.1.4/30",
+                "1.1.1.8/29",
+                "1.1.1.16/28",
+                "1.1.1.32/32",
+                "0:0:0:0:0:1:ffff:ffff/128",
+                "0:0:0:0:0:2:0:0/122",
+                "0:0:0:0:0:2:0:40/124",
+                "0:0:0:0:0:2:0:50/125",
+                "0:0:0:0:0:2:0:58/127",
+            },
+            true
+        ),
+        new networkTest(
+            "1.1.1.16",
+            28,
+            "mixed",
+            new String[]{
+                "1.1.1.16/28"
+            }
+        ),
+        new networkTest(
+            "1.1.1.4",
+            30,
+            "ipv4",
+            new String[]{
+                "1.1.1.4/30"
+            }
+        )
     };
 
     @Test