field-list flag to support self-define fields for the region

maker/java/src/main/java/org/lionsoul/ip2region/MakerTest.java

@@ -11,6 +11,10 @@ import org.lionsoul.ip2region.xdb.IndexPolicy;
 import org.lionsoul.ip2region.xdb.Log;
 import org.lionsoul.ip2region.xdb.Maker;
 
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+
 public class MakerTest {
 
     public final static Log log = Log.getLogger(MakerTest.class);
@@ -21,11 +25,55 @@ public class MakerTest {
         System.out.println("options:");
         System.out.println(" --src string        source ip text file path");
         System.out.println(" --dst string        destination binary xdb file path");
+        System.out.println(" --field-list string field index list imploded with ',' eg: 0,1,2,3,4");
         System.out.println(" --log-level string  set the log level, options: debug/info/warn/error");
     }
 
+    private static int[] getFieldList(String fieldList) {
+        final ArrayList<Integer> list = new ArrayList<Integer>();
+        final Map<String, String> map = new HashMap<String, String>();
+        if (!fieldList.isEmpty()) {
+            String[] fList =  fieldList.split(",");
+            for (String f : fList) {
+                final String s = f.trim();
+                if (s.isEmpty()) {
+                    log.infof("undefined option `%s`", f);
+                    return null;
+                }
+
+                if (!s.matches("^\\d+$")) {
+                    log.infof("field `%s` is not a number", f);
+                    return null;
+                }
+
+                if (map.containsKey(s)) {
+                    log.infof("duplicate field index `%s`", s);
+                    return null;
+                }
+
+                map.put(s, s);
+                final int idx = Integer.parseInt(s);
+                if (idx < 0) {
+                    log.infof("field index `%s` is negative", s);
+                    return null;
+                }
+
+                list.add(idx);
+            }
+        }
+
+        // let's keep it a complex way so the old JDK could run these pieces.
+        final int[] fields = new int[list.size()];
+        for (int i = 0; i < list.size(); i++) {
+            fields[i] = list.get(i);
+        }
+
+        return fields;
+    }
+
     public static void genDb(String[] args) throws Exception {
-        String srcFile = "", dstFile = "", logLevel = "info";
+        String srcFile = "", dstFile = "";
+        String fieldList = "", logLevel = "info";
         int indexPolicy = IndexPolicy.Vector;
         for (final String r : args) {
             if (r.length() < 5) {
@@ -49,6 +97,8 @@ public class MakerTest {
                 srcFile = val;
             } else if ("dst".equals(key)) {
                 dstFile = val;
+            } else if ("field-list".equals(key)) {
+                    fieldList = val;
             } else if ("log-level".equals(key)) {
                 logLevel = val;
             } else {
@@ -62,8 +112,14 @@ public class MakerTest {
             return;
         }
 
+        final int[] fields = getFieldList(fieldList);
+        if (fields == null) {
+            return;
+        }
+
+        // check and make the field list
         long tStart = System.currentTimeMillis();
-        final Maker maker = new Maker(indexPolicy, srcFile, dstFile);
+        final Maker maker = new Maker(indexPolicy, srcFile, dstFile, fields);
         log.infof("Generating xdb with src=%s, dst=%s, logLevel=%s\n", srcFile, dstFile, logLevel);
         Maker.log.setLevel(logLevel);
         maker.init();

maker/java/src/main/java/org/lionsoul/ip2region/UtilTest.java

@@ -58,6 +58,15 @@ public class UtilTest {
         }
     }
 
+    public static void testFieldFilter() {
+        final String reg = "0|16777215|保留|保留|保留|保留|||||||||||";
+        final String[] fs = reg.split("\\|", -1);
+        log.infof("length: %d", fs.length);
+        for (String s : fs) {
+            log.infof("%s", s);
+        }
+    }
+
     public static void main(String[] args) {
         System.out.println("+-- testing index policy");
         testIndexPolicy();
@@ -70,6 +79,10 @@ public class UtilTest {
         System.out.println("+-- testing segments");
         testSegment();
         System.out.println();
+
+        System.out.println("+-- testing field filter");
+        testFieldFilter();
+        System.out.println();
     }
 
 }

maker/java/src/main/java/org/lionsoul/ip2region/xdb/Maker.java

@@ -53,10 +53,7 @@ package org.lionsoul.ip2region.xdb;
 
 import java.io.*;
 import java.nio.charset.Charset;
-import java.util.HashMap;
+import java.util.*;
-import java.util.LinkedList;
-import java.util.List;
-import java.util.Map;
 
 public class Maker {
     // constants define
@@ -72,6 +69,7 @@ public class Maker {
 
     // source text file handle
     private final File srcFile;
+    private final int[] fields;
     private final List<Segment> segments;
     private final Charset bytesCharset;
 
@@ -87,12 +85,14 @@ public class Maker {
     // vector index raw bytes
     private final byte[] vectorIndex;
 
-    public Maker(int policy, String srcPath, String dstPath) throws IOException {
+    public Maker(int policy, String srcPath, String dstPath, int[] fields) throws IOException {
         this.srcFile = new File(srcPath);
         if (!this.srcFile.exists()) {
             throw new FileNotFoundException("source text file `" + srcPath + "` not found");
         }
 
+        this.fields = fields;
+
         /// check and delete the target xdb file if it exists
         /// final File dstFile = new File(dstPath);
         /// if (dstFile.exists() && !dstFile.delete()) {
@@ -128,6 +128,31 @@ public class Maker {
         dstHandle.write(header);
     }
 
+    // internal method to apply the region fields filter
+    private String getFilteredRegion(String region) {
+        if (this.fields.length == 0) {
+            return region;
+        }
+
+        final String[] fs = region.split("\\|", -1);
+        final StringBuilder sb = new StringBuilder();
+        final int tailing = this.fields.length - 1;
+        for (int i = 0; i < this.fields.length; i++) {
+            final int idx = this.fields[i];
+            if (idx >= fs.length) {
+                throw new IllegalArgumentException("field index `"
+                        + idx + "` exceeded the max length `" + fs.length + "`");
+            }
+
+            sb.append(fs[idx]);
+            if (sb.length() > 0 && i < tailing) {
+                sb.append("|");
+            }
+        }
+
+        return sb.toString();
+    }
+
     // load all the segments
     private void loadSegments() throws Exception {
         log.infof("try to load the segments ... ");
@@ -152,7 +177,7 @@ public class Maker {
                 throw new Exception("start ip("+ps[0]+") should not be greater than end ip("+ps[1]+")");
             }
 
-            if (ps[2].length() < 1) {
+            if (ps[2].isEmpty()) {
                 br.close();
                 throw new Exception("empty region info in segment line `"+ps[2]+"`");
             }
@@ -165,7 +190,7 @@ public class Maker {
                 }
             }
 
-            final Segment seg = new Segment(sip, eip, ps[2]);
+            final Segment seg = new Segment(sip, eip, getFilteredRegion(ps[2]));
             segments.add(seg);
             last = seg;
         }