SIMD strWidth ASCII optimization

2 files changed, 44 insertions, 5 deletions

diff --git a/src/display_width.zig b/src/display_width.zig
index e06aa8f..b96f872 100644
--- a/src/display_width.zig
+++ b/src/display_width.zig
@@ -1,4 +1,5 @@
 const std = @import("std");
+const simd = std.simd;
 const testing = std.testing;
 
 const CodePointIterator = @import("CodePoint").CodePointIterator;
@@ -14,8 +15,42 @@ pub fn codePointWidth(cp: u21) i3 {
     return dwp.stage_2[dwp.stage_1[cp >> 8] + (cp & 0xff)];
 }
 
-fn strWidth(str: []const u8) usize {
+fn isAsciiOnly(str: []const u8) bool {
+    const vec_len = simd.suggestVectorLength(u8) orelse @panic("No SIMD support.");
+    const Vec = @Vector(vec_len, u8);
+    var i: usize = 0;
+
+    while (i < str.len) : (i += vec_len) {
+        if (str[i..].len < vec_len) return for (str[i..]) |b| {
+            if (b > 127) break false;
+        } else true;
+
+        const v1 = str[i..].ptr[0..vec_len].*;
+        const v2: Vec = @splat(127);
+        if (@reduce(.Or, v1 > v2)) return false;
+    }
+
+    return true;
+}
+
+/// strWidth returns the total display width of `str` as the number of cells
+/// required in a fixed-pitch font (i.e. a terminal screen).
+pub fn strWidth(str: []const u8) usize {
     var total: isize = 0;
+
+    if (isAsciiOnly(str)) {
+        for (str) |b| {
+            // Backspace and delete
+            if (b == 0x8 or b == 0x7f) {
+                total -= 1;
+            } else if (b >= 0x20) {
+                total += 1;
+            }
+        }
+
+        return if (total > 0) @intCast(total) else 0;
+    }
+
     var giter = GraphemeIterator.init(str);
 
     while (giter.next()) |gc| {
diff --git a/src/main.zig b/src/main.zig
index 3e65c7b..e7c0828 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -3,7 +3,9 @@ const std = @import("std");
 // const GraphemeIterator = @import("ziglyph").GraphemeIterator;
 // const GraphemeIterator = @import("Grapheme").GraphemeIterator;
 // const codePointWidth = @import("ziglyph").display_width.codePointWidth;
-const codePointWidth = @import("display_width").codePointWidth;
+// const codePointWidth = @import("display_width").codePointWidth;
+// const strWidth = @import("ziglyph").display_width.strWidth;
+const strWidth = @import("display_width").strWidth;
 const CodePointIterator = @import("CodePoint").CodePointIterator;
 
 pub fn main() !void {
@@ -14,14 +16,16 @@ pub fn main() !void {
     const input = try std.fs.cwd().readFileAlloc(allocator, "lang_mix.txt", std.math.maxInt(u32));
     defer allocator.free(input);
 
-    var result: isize = 0;
+    var result: usize = 0;
     // var iter = GraphemeIterator.init(input);
-    var iter = CodePointIterator{ .bytes = input };
+    // var iter = CodePointIterator{ .bytes = input };
+    var iter = std.mem.splitScalar(u8, input, '\n');
 
     var timer = try std.time.Timer.start();
 
     // for (0..50) |_| {
-    while (iter.next()) |cp| result += codePointWidth(@intCast(cp.code));
+    // while (iter.next()) |cp| result += codePointWidth(@intCast(cp.code));
+    while (iter.next()) |line| result += strWidth(line);
     // iter.cp_iter.i = 0;
     // }