From f3115ce8640ea314c35b391f84e26c7bd599fc0e Mon Sep 17 00:00:00 2001
From: Jean Boussier <jean.boussier@gmail.com>
Date: Wed, 27 Nov 2024 10:27:09 +0100
Subject: [PATCH] Micro optimize encoding checks
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Profiling shows a lot of time spent in various encoding check functions.
I'm working on optimizing them on the Ruby side, but if we assume most
strings are one of the simple 3 encodings, we can skip a lot of overhead.

```ruby
require 'strscan'
require 'benchmark/ips'

source = 10_000.times.map { rand(9999999).to_s }.join(",").force_encoding(Encoding::UTF_8).freeze

def scan_to_i(source)
  scanner = StringScanner.new(source)
  while number = scanner.scan(/\d+/)
    number.to_i
    scanner.skip(",")
  end
end

def scan_integer(source)
  scanner = StringScanner.new(source)
  while scanner.scan_integer
    scanner.skip(",")
  end
end

Benchmark.ips do |x|
  x.report("scan.to_i") { scan_to_i(source) }
  x.report("scan_integer") { scan_integer(source) }
  x.compare!
end
```

Before:

```
ruby 3.3.4 (2024-07-09 revision be1089c8ec) +YJIT [arm64-darwin23]
Warming up --------------------------------------
           scan.to_i    93.000 i/100ms
        scan_integer   232.000 i/100ms
Calculating -------------------------------------
           scan.to_i    933.191 (± 0.2%) i/s    (1.07 ms/i) -      4.743k in   5.082597s
        scan_integer      2.326k (± 0.8%) i/s  (429.99 μs/i) -     11.832k in   5.087974s

Comparison:
        scan_integer:     2325.6 i/s
           scan.to_i:      933.2 i/s - 2.49x  slower
```

After:

```
ruby 3.3.4 (2024-07-09 revision be1089c8ec) +YJIT [arm64-darwin23]
Warming up --------------------------------------
           scan.to_i    96.000 i/100ms
        scan_integer   274.000 i/100ms
Calculating -------------------------------------
           scan.to_i    969.489 (± 0.2%) i/s    (1.03 ms/i) -      4.896k in   5.050114s
        scan_integer      2.756k (± 0.1%) i/s  (362.88 μs/i) -     13.974k in   5.070837s

Comparison:
        scan_integer:     2755.8 i/s
           scan.to_i:      969.5 i/s - 2.84x  slower
```
---
 ext/strscan/strscan.c | 41 ++++++++++++++++++++++++++++++++++++++---
 1 file changed, 38 insertions(+), 3 deletions(-)

diff --git a/ext/strscan/strscan.c b/ext/strscan/strscan.c
index 165cd4271b..51ffe5af69 100644
--- a/ext/strscan/strscan.c
+++ b/ext/strscan/strscan.c
@@ -32,6 +32,8 @@ static VALUE StringScanner;
 static VALUE ScanError;
 static ID id_byteslice;
 
+static int usascii_encindex, utf8_encindex, binary_encindex;
+
 struct strscanner
 {
     /* multi-purpose flags */
@@ -683,6 +685,14 @@ strscan_search(regex_t *reg, VALUE str, struct re_registers *regs, void *args_pt
                        ONIG_OPTION_NONE);
 }
 
+static void
+strscan_enc_check(VALUE str1, VALUE str2)
+{
+    if (RB_ENCODING_GET(str1) != RB_ENCODING_GET(str2)) {
+        rb_enc_check(str1, str2);
+    }
+}
+
 static VALUE
 strscan_do_scan(VALUE self, VALUE pattern, int succptr, int getstr, int headonly)
 {
@@ -710,18 +720,21 @@ strscan_do_scan(VALUE self, VALUE pattern, int succptr, int getstr, int headonly
     }
     else {
         StringValue(pattern);
-        rb_encoding *enc = rb_enc_check(p->str, pattern);
         if (S_RESTLEN(p) < RSTRING_LEN(pattern)) {
+            strscan_enc_check(p->str, pattern);
             return Qnil;
         }
 
         if (headonly) {
+            strscan_enc_check(p->str, pattern);
+
             if (memcmp(CURPTR(p), RSTRING_PTR(pattern), RSTRING_LEN(pattern)) != 0) {
                 return Qnil;
             }
             set_registers(p, RSTRING_LEN(pattern));
         }
         else {
+            rb_encoding *enc = rb_enc_check(p->str, pattern);
             long pos = rb_memsearch(RSTRING_PTR(pattern), RSTRING_LEN(pattern),
                                     CURPTR(p), S_RESTLEN(p), enc);
             if (pos == -1) {
@@ -1282,6 +1295,24 @@ strscan_parse_integer(struct strscanner *p, int base, long len)
     return integer;
 }
 
+static inline bool
+strscan_ascii_compat_fastpath(VALUE str) {
+    int encindex = ENCODING_GET_INLINED(str);
+    // The overwhelming majority of strings are in one of these 3 encodings.
+    return encindex == utf8_encindex || encindex == binary_encindex || encindex == usascii_encindex;
+}
+
+static inline void
+strscan_must_ascii_compat(VALUE str)
+{
+    // The overwhelming majority of strings are in one of these 3 encodings.
+    if (RB_LIKELY(strscan_ascii_compat_fastpath(str))) {
+        return;
+    }
+
+    rb_must_asciicompat(str);
+}
+
 static VALUE
 strscan_scan_base10_integer(VALUE self)
 {
@@ -1292,7 +1323,7 @@ strscan_scan_base10_integer(VALUE self)
     GET_SCANNER(self, p);
     CLEAR_MATCH_STATUS(p);
 
-    rb_must_asciicompat(p->str);
+    strscan_must_ascii_compat(p->str);
 
     ptr = CURPTR(p);
 
@@ -1330,7 +1361,7 @@ strscan_scan_base16_integer(VALUE self)
     GET_SCANNER(self, p);
     CLEAR_MATCH_STATUS(p);
 
-    rb_must_asciicompat(p->str);
+    strscan_must_ascii_compat(p->str);
 
     ptr = CURPTR(p);
 
@@ -2251,6 +2282,10 @@ Init_strscan(void)
 
     id_byteslice = rb_intern("byteslice");
 
+    usascii_encindex = rb_usascii_encindex();
+    utf8_encindex = rb_utf8_encindex();
+    binary_encindex = rb_ascii8bit_encindex();
+
     StringScanner = rb_define_class("StringScanner", rb_cObject);
     ScanError = rb_define_class_under(StringScanner, "Error", rb_eStandardError);
     if (!rb_const_defined(rb_cObject, id_scanerr)) {