Updated the vectorized scanner to more closely match the C implmeentation.

Author: samyron

java/src/json/ext/EscapeScanner.java

@@ -15,33 +15,42 @@ static class State {
     }
 
     static class VectorSupport {
-        static Constructor<?> vectorizedEscapeScannerConstructor = null;
+        static final EscapeScanner VECTORIZED_ESCAPE_SCANNER;
 
         static {
             Optional<Module> vectorModule = ModuleLayer.boot().findModule("jdk.incubator.vector");
+            EscapeScanner scanner = null;
             if (vectorModule.isPresent()) {
                 try {
                     Class<?> vectorEscapeScannerClass = EscapeScanner.class.getClassLoader().loadClass("json.ext.VectorizedEscapeScanner");
-                    vectorizedEscapeScannerConstructor = vectorEscapeScannerClass.getDeclaredConstructor();
-                } catch (ClassNotFoundException | NoSuchMethodException e) {
+                    Constructor<?> vectorizedEscapeScannerConstructor = vectorEscapeScannerClass.getDeclaredConstructor();
+                    scanner = (EscapeScanner) vectorizedEscapeScannerConstructor.newInstance();
+                } catch (ClassNotFoundException | NoSuchMethodException | InstantiationException | IllegalAccessException | InvocationTargetException e) {
                     // Fallback to the ScalarEscapeScanner if we cannot load the VectorizedEscapeScanner.
                     System.err.println("Failed to load VectorizedEscapeScanner, falling back to ScalarEscapeScanner: " + e.getMessage());
+                    scanner = null;
                 }
+                
             }
+            VECTORIZED_ESCAPE_SCANNER = scanner;
         }
     }
 
     boolean scan(EscapeScanner.State state) throws java.io.IOException;
 
-    public static EscapeScanner basicScanner() {
-        if (VectorSupport.vectorizedEscapeScannerConstructor != null) {
-            try {
-                // Attempt to instantiate the vectorized escape scanner if available.
-                return (EscapeScanner) VectorSupport.vectorizedEscapeScannerConstructor.newInstance();
-            } catch (InstantiationException | IllegalAccessException | InvocationTargetException e) {
-                System.err.println("Failed to instantiate VectorizedEscapeScanner, falling back to ScalarEscapeScanner: " + e.getMessage());
-            }
+    default State createState(byte[] ptrBytes, int ptr, int len, int beg) {
+        State state = new State();
+        state.ptrBytes = ptrBytes;
+        state.ptr = ptr;
+        state.len = len;
+        state.beg = beg;
+        state.pos = 0; // Start scanning from the beginning of the segment
+        return state;
+    }
 
+    public static EscapeScanner basicScanner() {
+        if (VectorSupport.VECTORIZED_ESCAPE_SCANNER != null) {
+            return VectorSupport.VECTORIZED_ESCAPE_SCANNER;
         }
 
         return new ScalarEscapeScanner(StringEncoder.ESCAPE_TABLE);

java/src/json/ext/StringEncoder.java

@@ -228,17 +228,18 @@ boolean searchEscape(EscapeScanner.State state) throws IOException {
     }
 
     void encodeBasic(ByteList src) throws IOException {
-        EscapeScanner.State state = new EscapeScanner.State();
-        state.ptrBytes = src.unsafeBytes();
-        state.ptr = src.begin();
-        state.len = src.realSize();
-        state.beg = 0;
-        state.pos = 0;
+        // EscapeScanner.State state = new EscapeScanner.State();
+        // state.ptrBytes = src.unsafeBytes();
+        // state.ptr = src.begin();
+        // state.len = src.realSize();
+        // state.beg = 0;
+        // state.pos = 0;
 
         byte[] hexdig = HEX;
         byte[] scratch = aux;
 
         EscapeScanner scanner = EscapeScanner.basicScanner();
+        EscapeScanner.State state = scanner.createState(src.unsafeBytes(), src.begin(), src.realSize(), 0);
 
         while(scanner.scan(state)) {
             int ch = Byte.toUnsignedInt(state.ptrBytes[state.ptr + state.pos]);

java/src/json/ext/VectorizedEscapeScanner.java

@@ -1,32 +1,31 @@
 package json.ext;
 
 import java.io.IOException;
+import javax.naming.directory.NoSuchAttributeException;
 
 import jdk.incubator.vector.ByteVector;
 import jdk.incubator.vector.VectorMask;
 import jdk.incubator.vector.VectorOperators;
 import jdk.incubator.vector.VectorSpecies;
+import jdk.jfr.RecordingState;
 
 public class VectorizedEscapeScanner implements EscapeScanner {
     public static EscapeScanner.ScalarEscapeScanner FALLBACK = new EscapeScanner.ScalarEscapeScanner(StringEncoder.ESCAPE_TABLE);
 
-    // private VectorMask<Byte> needsEscape = null;
-    // private int chunkStart = 0;
-
     @Override
-    public boolean scan(State state) throws IOException {
+    public boolean scan(State _state) throws IOException {
         VectorSpecies<Byte> species = ByteVector.SPECIES_PREFERRED;
 
-        // if (needsEscape != null) {
-        //     if (needsEscape.anyTrue()) {
-        //         int firstEscapeIndex = needsEscape.firstTrue();
-        //         needsEscape = needsEscape.andNot(VectorMask.fromLong(species, 1L << firstEscapeIndex));
-        //         state.pos = chunkStart + firstEscapeIndex;
-        //         return true;
-        //     } else {
-        //         needsEscape = null;
-        //     }
-        // }
+        VectorizedState state = (VectorizedState) _state;
+
+        if (state.hasMatches) {
+            if (state.mask > 0) {
+                return nextMatch(state);
+            } else {
+                state.hasMatches = false;
+                state.pos = state.chunkStart + species.length();
+            }
+        }
 
         while ((state.ptr + state.pos) + species.length() < state.len) {
             ByteVector chunk = ByteVector.fromArray(species, state.ptrBytes, state.ptr + state.pos);
@@ -41,17 +40,41 @@ public boolean scan(State state) throws IOException {
 
             VectorMask<Byte> needsEscape = chunk.eq(ByteVector.broadcast(species, '\\')).or(tooLowOrDblQuote).and(negative);
             if (needsEscape.anyTrue()) {
-                // chunkStart = state.ptr + state.pos;
-                int firstEscapeIndex = needsEscape.firstTrue();
-                // Clear the bit at firstEscapeIndex to avoid scanning the same byte again
-                // needsEscape = needsEscape.andNot(VectorMask.fromLong(species, 1L << firstEscapeIndex));
-                state.pos += firstEscapeIndex;
-                return true;
+                state.hasMatches = true;
+                state.chunkStart = state.ptr + state.pos;
+                state.mask = needsEscape.toLong();
+
+                return nextMatch(state);
             }
 
             state.pos += species.length();
         }
 
         return FALLBACK.scan(state);
     }
+
+    private boolean nextMatch(VectorizedState state) {
+        int index = Long.numberOfTrailingZeros(state.mask);
+        state.mask &= (state.mask - 1);
+        state.pos = state.chunkStart + index;
+        return true;
+    }
+
+    @Override
+    public EscapeScanner.State createState(byte[] ptrBytes, int ptr, int len, int beg) {
+        VectorizedState state = new VectorizedState();
+        state.ptrBytes = ptrBytes;
+        state.ptr = ptr;
+        state.len = len;
+        state.beg = beg;
+        state.pos = 0; 
+        return state;
+    }
+
+    private static class VectorizedState extends State {
+        private long mask;
+        private int chunkStart = 0;
+        // private int lastMatchingIndex;
+        private boolean hasMatches;
+    }
 }