usearch 2.9.0 → 2.9.2

package/stringzilla/cli/split.py

@@ -1,50 +1,91 @@
 #!/usr/bin/env python3
 
+import argparse
 import sys
-
+import stringzilla
 from stringzilla import File, Str
 
 
-def split_file(file_path, lines_per_file, output_prefix):
+def parse_arguments():
+    parser = argparse.ArgumentParser(
+        description="Output pieces of FILE to PREFIXaa, PREFIXab, ...; default size is 1000 lines, and default PREFIX is 'x'."
+    )
+    parser.add_argument(
+        "file", nargs="?", default="-", help='File to process, "-" for standard input'
+    )
+    parser.add_argument(
+        "prefix", nargs="?", default="X", help='Output file prefix, default is "x"'
+    )
+    parser.add_argument(
+        "-l",
+        "--lines",
+        type=int,
+        default=1000,
+        help="Number of lines per output file, default is 1000",
+    )
+    parser.add_argument(
+        "-t",
+        "--separator",
+        default="\n",
+        help="Use SEP instead of newline as the record separator; '\\0' (zero) specifies the NUL character",
+    )
+    parser.add_argument(
+        "-n",
+        "--number",
+        type=int,
+        default=None,
+        help="Generate N output files based on size of input",
+    )
+    parser.add_argument("--version", action="version", version=stringzilla.__version__)
+    return parser.parse_args()
+
+
+def split_file(file_path, lines_per_file, output_prefix, separator, number_of_files):
     try:
-        # 1. Memory-map the large file
-        file_mapped = File(file_path)
-        file_contents = Str(file_mapped)
-
-        # Variables to keep track of the current position and file part number
+        if separator == "\\0":
+            separator = "\0"
+        if file_path == "-":
+            file_contents = Str(sys.stdin.read())
+        else:
+            file_mapped = File(file_path)
+            file_contents = Str(file_mapped)
+
+        if number_of_files is not None:
+            total_length = len(file_contents)
+            chunk_size = total_length // number_of_files
+            for file_part in range(number_of_files):
+                start = file_part * chunk_size
+                end = (
+                    start + chunk_size
+                    if file_part < number_of_files - 1
+                    else total_length
+                )
+                current_slice = file_contents[start:end]
+                output_path = f"{output_prefix}{file_part}"
+                current_slice.write_to(output_path)
+            return
         current_position = 0
         file_part = 0
-        newline_position = (
-            -1
-        )  # Start before file begins to find the first newline correctly
+        newline_position = -1
 
-        # Loop until the end of the file
         while current_position < len(file_contents):
-            # 2. Loop to skip `lines_per_file` lines
             for _ in range(lines_per_file):
-                newline_position = file_contents.find("\n", newline_position + 1)
-                if newline_position == -1:  # No more newlines
+                newline_position = file_contents.find(separator, newline_position + 1)
+                if newline_position == -1:
                     break
 
-            # If no newlines were found and we're not at the start, process the rest of the file
             if newline_position == -1 and current_position < len(file_contents):
                 newline_position = len(file_contents)
 
-            # 3. Use offset_within to get the length of the current section
-            # Assuming offset_within gives you the length from the current position
             section_length = (
                 newline_position - current_position if newline_position != -1 else 0
             )
 
-            # Extract the current section to write out
-            if section_length > 0:  # Prevent creating empty files
+            if section_length > 0:
                 current_slice = file_contents[current_position : newline_position + 1]
-
-                # 4. Save the current slice to file
                 output_path = f"{output_prefix}{file_part}"
                 current_slice.write_to(output_path)
 
-                # Prepare for the next slice
                 file_part += 1
                 current_position = newline_position + 1
 
@@ -52,20 +93,12 @@ def split_file(file_path, lines_per_file, output_prefix):
         print(f"No such file: {file_path}")
     except Exception as e:
         print(f"An error occurred: {e}")
+        print("Usage example: split.py [-l LINES] [file] [prefix]")
 
 
 def main():
-    if len(sys.argv) < 4:
-        print(
-            "Usage: python split_file.py <lines_per_file> <input_file> <output_prefix>"
-        )
-        sys.exit(1)
-
-    lines_per_file = int(sys.argv[1])
-    file_path = sys.argv[2]
-    output_prefix = sys.argv[3]
-
-    split_file(file_path, lines_per_file, output_prefix)
+    args = parse_arguments()
+    split_file(args.file, args.lines, args.prefix, args.separator, args.number)
 
 
 if __name__ == "__main__":

package/stringzilla/cli/wc.py

@@ -1,36 +1,116 @@
 #!/usr/bin/env python3
 
 import sys
-
+import argparse
+import stringzilla
 from stringzilla import File, Str
 
 
-def wc(file_path):
-    try:
-        mapped_file = File(file_path)
-        mapped_bytes = Str(mapped_file)
-        line_count = mapped_bytes.count("\n")
-        word_count = mapped_bytes.count(" ")
-        char_count = mapped_bytes.__len__()
+def parse_arguments():
+    parser = argparse.ArgumentParser(
+        description="Print newline, word, and byte counts for each FILE, and a total line if more than one FILE is \
+        specified. A word is a non-zero-length sequence of acters delimited by white space."
+    )
+    parser.add_argument("files", nargs="*", default=["-"], help="Files to process")
+    parser.add_argument(
+        "-c", "--bytes", action="store_true", help="print the byte counts"
+    )
+    parser.add_argument(
+        "-m", "--chars", action="store_true", help="print the character counts"
+    )
+    parser.add_argument(
+        "-l", "--lines", action="store_true", help="print the newline counts"
+    )
+    parser.add_argument(
+        "-L",
+        "--max-line-length",
+        action="store_true",
+        help="print the maximum display width",
+    )
+    parser.add_argument(
+        "-w", "--words", action="store_true", help="print the word counts"
+    )
+    parser.add_argument("--version", action="version", version=stringzilla.__version__)
+    return parser.parse_args()
+
+
+def wc(file_path, args):
+    if file_path == "-":  # read from stdin
+        content = sys.stdin.read()
+        mapped_bytes = Str(content)
+    else:
+        try:
+            mapped_file = File(file_path)
+            mapped_bytes = Str(mapped_file)
+        except RuntimeError:  # File gives a RuntimeError if the file does not exist
+            return f"No such file: {file_path}", False
+
+    line_count = mapped_bytes.count("\n")
+    word_count = mapped_bytes.count(" ") + 1
+    char_count = mapped_bytes.__len__()
+    counts = {
+        "line_count": line_count,
+        "word_count": word_count,
+        "char_count": char_count,
+    }
+
+    if args.max_line_length:
+        max_line_length = max(len(line) for line in str(mapped_bytes).split("\n"))
+        counts["max_line_length"] = max_line_length
 
-        return line_count, word_count, char_count
-    except FileNotFoundError:
-        return f"No such file: {file_path}"
+    if args.bytes or args.chars:
+        byte_count = char_count  # assume 1 char = 1 byte
+        counts["byte_count"] = byte_count
+
+    return counts, True
+
+
+def format_output(counts, args):
+    selected_counts = []
+    if args.lines:
+        selected_counts.append(counts["line_count"])
+    if args.words:
+        selected_counts.append(counts["word_count"])
+    if args.chars:
+        selected_counts.append(counts["char_count"])
+    if args.bytes:
+        selected_counts.append(counts.get("byte_count", counts["char_count"]))
+    if args.max_line_length:
+        selected_counts.append(counts.get("max_line_length", 0))
+
+    if not any([args.lines, args.words, args.chars, args.bytes, args.max_line_length]):
+        selected_counts = [
+            counts["line_count"],
+            counts["word_count"],
+            counts["char_count"],
+        ]
+
+    return " ".join(str(count) for count in selected_counts)
 
 
 def main():
-    if len(sys.argv) < 2:
-        print("Usage: python wc.py <file>")
-        sys.exit(1)
+    args = parse_arguments()
+    total_counts = {
+        "line_count": 0,
+        "word_count": 0,
+        "char_count": 0,
+        "max_line_length": 0,
+        "byte_count": 0,
+    }
 
-    file_path = sys.argv[1]
-    counts = wc(file_path)
+    for file_path in args.files:
+        counts, success = wc(file_path, args)
+        if success:
+            for key in total_counts.keys():
+                total_counts[key] += counts.get(key, 0)
+            output = format_output(counts, args) + f" {file_path}"
+            print(output)
+        else:
+            print(counts)
 
-    if isinstance(counts, tuple):
-        line_count, word_count, char_count = counts
-        print(f"{line_count} {word_count} {char_count} {file_path}")
-    else:
-        print(counts)
+    if len(args.files) > 1:
+        total_output = format_output(total_counts, args) + " total"
+        print(total_output)
 
 
 if __name__ == "__main__":

package/stringzilla/include/stringzilla/experimental.h

@@ -368,6 +368,127 @@ sz_u512_vec_t sz_inclusive_min(sz_i32_t previous, sz_error_cost_t gap, sz_u512_v
 
 #if SZ_USE_ARM_NEON
 
+SZ_PUBLIC sz_cptr_t sz_find_neon_too_smart(sz_cptr_t h, sz_size_t h_length, sz_cptr_t n, sz_size_t n_length) {
+
+    // This almost never fires, but it's better to be safe than sorry.
+    if (h_length < n_length || !n_length) return SZ_NULL_CHAR;
+    if (n_length == 1) return sz_find_byte_neon(h, h_length, n);
+
+    // Scan through the string.
+    // Assuming how tiny the Arm NEON registers are, we should avoid internal branches at all costs.
+    // That's why, for smaller needles, we use different loops.
+    if (n_length == 2) {
+        // This is a common case. Aside from ASCII bigrams, it's also the most common case for UTF-16,
+        // or any UTF8 content in Cyrillic, Greek, Armenian, Hebrew, Arabic, Coptic, Syriac, Thaana,
+        // N'Ko writing system of West-African nations, and, of course, Latin scripts.
+        // Dealing with 16-bit values, we can check 15 possible offsets in a single loop iteration.
+        // For that we are going to keep 2 registers populated with haystack data.
+        // First - bigrams at even offsets - 0, 2, 4, 6, 8, 10, 12, 14.
+        // Second - bigrams at odd offsets - 1, 3, 5, 7, 9, 11, 13. One less than the first one.
+        sz_u64_t matches;
+        sz_u128_vec_t h_even_vec, h_odd_vec, n_vec, interleave_mask_vec, matches_vec;
+        // Broadcast needle characters into SIMD registers.
+        n_vec.u16x8 = vdupq_n_u16(sz_u16_load(n).u16);
+        interleave_mask_vec.u16x8 = vdupq_n_u16(0x00FFu);
+        for (; h_length >= 16; h += 15, h_length -= 15) {
+            h_even_vec.u8x16 = vld1q_u8((sz_u8_t const *)h);
+            h_odd_vec.u8x16 = vextq_u8(h_even_vec.u8x16, /* can be any noise: */ h_even_vec.u8x16, 1);
+            // We can now compare both 16-bit arrays with the needle.
+            // The result of each comparison will also be 16 bits long.
+            // Then - we blend!
+            // For odd offsets we are gonna take the bottom 8 bits, and for even - the top ones!
+            matches_vec.u8x16 =
+                vbslq_u8(interleave_mask_vec.u8x16, vreinterpretq_u8_u16(vceqq_u16(h_even_vec.u16x8, n_vec.u16x8)),
+                         vreinterpretq_u8_u16(vceqq_u16(h_odd_vec.u16x8, n_vec.u16x8)));
+            matches = vreinterpretq_u8_u4(matches_vec.u8x16);
+            if (matches) return h + sz_u64_ctz(matches) / 4;
+        }
+    }
+    else if (n_length == 3) {
+        // Comparing 24-bit values is a bumer. Being lazy, I went with a simple design.
+        // Instead of keeping one register per haystack offset, I keep a register per needle character.
+        sz_u64_t matches;
+        sz_u128_vec_t h_vec, n_first_vec, n_second_vec, n_third_vec, matches_vec;
+        // Broadcast needle characters into SIMD registers.
+        n_first_vec.u8x16 = vld1q_dup_u8((sz_u8_t const *)&n[0]);
+        n_second_vec.u8x16 = vld1q_dup_u8((sz_u8_t const *)&n[1]);
+        n_third_vec.u8x16 = vld1q_dup_u8((sz_u8_t const *)&n[2]);
+        for (; h_length >= 16; h += 14, h_length -= 14) {
+            h_vec.u8x16 = vld1q_u8((sz_u8_t const *)h);
+            // Let's compare the first character.
+            matches_vec.u8x16 = vceqq_u8(h_vec.u8x16, n_first_vec.u8x16);
+            // Let's compare the second one, shift the equality indicators left by 8 bits, and blend.
+            matches_vec.u8x16 =
+                vandq_u8(matches_vec.u8x16, vextq_u8(vceqq_u8(h_vec.u8x16, n_second_vec.u8x16), vdupq_n_u8(0), 1));
+            // Let's compare the third one, shift the equality indicators left by 16 bits, and blend.
+            matches_vec.u8x16 =
+                vandq_u8(matches_vec.u8x16, vextq_u8(vceqq_u8(h_vec.u8x16, n_third_vec.u8x16), vdupq_n_u8(0), 2));
+            // Now reduce bytes to nibbles, and check for matches.
+            matches = vreinterpretq_u8_u4(matches_vec.u8x16);
+            if (matches) return h + sz_u64_ctz(matches) / 4;
+        }
+    }
+    else if (n_length == 4) {
+        // This is a common case not only for ASCII 4-grams, but also UTF-32 content,
+        // emojis, Chinese, and many other east-Asian languages.
+        // Dealing with 32-bit values, we can analyze 13 offsets at once.
+        sz_u128_vec_t h_first_vec, h_second_vec, h_third_vec, h_fourth_vec, interleave_2mask_vec, interleave_4mask_vec,
+            n_vec, matches_vec;
+        sz_u64_t matches;
+        // Broadcast needle characters into SIMD registers.
+        n_vec.u32x4 = vdupq_n_u32(sz_u32_load(n).u32);
+        interleave_2mask_vec.u16x8 = vdupq_n_u16(0x00FFu);
+        interleave_4mask_vec.u32x4 = vdupq_n_u32(0x0000FFFFu);
+        for (; h_length >= 16; h += 13, h_length -= 13) {
+            h_first_vec.u8x16 = vld1q_u8((sz_u8_t const *)h);
+            h_second_vec.u8x16 = vextq_u8(h_first_vec.u8x16, /* can be any noise: */ h_first_vec.u8x16, 1);
+            h_third_vec.u8x16 = vextq_u8(h_first_vec.u8x16, /* can be any noise: */ h_first_vec.u8x16, 2);
+            h_fourth_vec.u8x16 = vextq_u8(h_first_vec.u8x16, /* can be any noise: */ h_first_vec.u8x16, 3);
+            // We can now compare all four arrays of 32-bit values with the needle.
+            // The result of each comparison will also be 32 bits long.
+            // Then - we blend!
+            matches_vec.u8x16 = vbslq_u8(
+                interleave_4mask_vec.u8x16,
+                vbslq_u8(interleave_2mask_vec.u8x16, vreinterpretq_u8_u32(vceqq_u32(h_first_vec.u32x4, n_vec.u32x4)),
+                         vreinterpretq_u8_u32(vceqq_u32(h_second_vec.u32x4, n_vec.u32x4))),
+                vbslq_u8(interleave_2mask_vec.u8x16, vreinterpretq_u8_u32(vceqq_u32(h_third_vec.u32x4, n_vec.u32x4)),
+                         vreinterpretq_u8_u32(vceqq_u32(h_fourth_vec.u32x4, n_vec.u32x4))));
+            matches = vreinterpretq_u8_u4(matches_vec.u8x16);
+            if (matches) return h + sz_u64_ctz(matches) / 4;
+        }
+    }
+    else {
+        // Pick the parts of the needle that are worth comparing.
+        sz_size_t offset_first, offset_mid, offset_last;
+        _sz_locate_needle_anomalies(n, n_length, &offset_first, &offset_mid, &offset_last);
+        // Broadcast those characters into SIMD registers.
+        sz_u64_t matches;
+        sz_u128_vec_t h_first_vec, h_mid_vec, h_last_vec, n_first_vec, n_mid_vec, n_last_vec, matches_vec;
+        n_first_vec.u8x16 = vld1q_dup_u8((sz_u8_t const *)&n[offset_first]);
+        n_mid_vec.u8x16 = vld1q_dup_u8((sz_u8_t const *)&n[offset_mid]);
+        n_last_vec.u8x16 = vld1q_dup_u8((sz_u8_t const *)&n[offset_last]);
+        // Walk through the string.
+        for (; h_length >= n_length + 16; h += 16, h_length -= 16) {
+            h_first_vec.u8x16 = vld1q_u8((sz_u8_t const *)(h + offset_first));
+            h_mid_vec.u8x16 = vld1q_u8((sz_u8_t const *)(h + offset_mid));
+            h_last_vec.u8x16 = vld1q_u8((sz_u8_t const *)(h + offset_last));
+            matches_vec.u8x16 = vandq_u8(                           //
+                vandq_u8(                                           //
+                    vceqq_u8(h_first_vec.u8x16, n_first_vec.u8x16), //
+                    vceqq_u8(h_mid_vec.u8x16, n_mid_vec.u8x16)),
+                vceqq_u8(h_last_vec.u8x16, n_last_vec.u8x16));
+            matches = vreinterpretq_u8_u4(matches_vec.u8x16);
+            while (matches) {
+                int potential_offset = sz_u64_ctz(matches) / 4;
+                if (sz_equal(h + potential_offset, n, n_length)) return h + potential_offset;
+                matches &= matches - 1;
+            }
+        }
+    }
+
+    return sz_find_serial(h, h_length, n, n_length);
+}
+
 SZ_INTERNAL void interleave_uint32x4_to_uint64x2(uint32x4_t in_low, uint32x4_t in_high, uint64x2_t *out_first_second,
                                                  uint64x2_t *out_third_fourth) {
     // Interleave elements
@@ -588,7 +709,7 @@ SZ_PUBLIC void sz_hashes_neon_reusing_loads(sz_cptr_t start, sz_size_t length, s
 }
 
 SZ_PUBLIC void sz_hashes_neon_readahead(sz_cptr_t start, sz_size_t length, sz_size_t window_length, sz_size_t step,
-                                       sz_hash_callback_t callback, void *callback_handle) {
+                                        sz_hash_callback_t callback, void *callback_handle) {
 
     if (length < window_length || !window_length) return;
     if (length < 2 * window_length) {