perf(typecheck): compare = guards over engine terms in the arena (b-typed ~4.3x less CPU)

CHANGELOG.md

@@ -37,6 +37,20 @@ Notable changes to shen-lua. Release notes for earlier versions are under
 - `bin/yggdrasil-build.lua` expects `kernel-version=42.2` manifests (older
   slices still build, with a warning).
 
+### Performance
+
+- Typechecking: the native Prolog translator compiles an `(= T1 T2)` guard
+  over two engine terms to an in-arena comparison (`E.eqterm`) instead of
+  copying both terms out to Shen lists and calling `=`. That guard is
+  `shen.by-hypothesis`'s check, run once per hypothesis per goal, and copying
+  it was ~75% of a typecheck-heavy run: capcov's `b-typed` workload drops
+  from ~48 s to ~11 s CPU (instructions retired 661 G -> 140 G) with
+  byte-identical output. Each side may be
+  `(shen.deref X B)` or a bare term variable, so the deref-hoisted
+  `by-hypothesis-h` of a patched 43.1 kernel takes the same path. Atoms and
+  variables that differ still go through the kernel's `=`, so tuples and
+  vectors keep structural equality.
+
 ### Fixed
 
 - Deep non-tail recursion no longer overflows the Lua stack, in user code

prolog_compile.lua

@@ -250,7 +250,32 @@ end
 -- into a KL defun called with materialized term arguments. Returns the Lua
 -- call expression.
 local guard_n = 0
+
+-- A guard argument that names one engine term: (shen.deref X B) or a bare
+-- term variable X (the driver may already have dereferenced X, as the 43.1
+-- deref-hoisted by-hypothesis-h does). Returns its Lua local, else nil.
+local function guard_term_ref(x, env)
+  if is_cons(x) and x[1] == SYM["shen.deref"] and is_cons(x[2]) then
+    x = x[2][1]
+  end
+  if is_klvar(x) then
+    local b = env[x.name]
+    if b and b.kind == "term" then return b.lua end
+  end
+  return nil
+end
+
 local function lift_guard(ctx, e, env, out, d)
+  -- (= T1 T2) over two engine terms: compare in the arena (E.eqterm) instead
+  -- of materializing both sides to Shen lists and calling `=`. This is the
+  -- by-hypothesis check, run once per hypothesis per goal; materializing it
+  -- was ~75% of a typecheck-heavy run.
+  if is_cons(e) and e[1] == SYM["="] and is_cons(e[2]) and is_cons(e[2][2])
+     and e[2][2][2] == NIL then
+    local a = guard_term_ref(e[2][1], env)
+    local b = a and guard_term_ref(e[2][2][1], env)
+    if b then return "EQTERM(" .. a .. ", " .. b .. ")" end
+  end
   -- replace (shen.deref X B) / (shen.lazyderef X B) with fresh params
   local params, args = {}, {}
   local seen = {}   -- engine var -> param (dedupe repeated references)
@@ -976,6 +1001,7 @@ local E, NP, F, MAT, IMP = ...
 local CONS, CAR, CDR, LZD = E.cons, E.car, E.cdr, E.lazyderef
 local NEWVAR, POPVAR = E.newvar, E.popvar
 local UNIFY, UNIFYOC, BIND1 = E.unify, E.unify_oc, E.bind1
+local EQTERM = E.eqterm
 local CUT, UNLOCK, LOCKOPEN = E.cut, E.unlock, E.lock_is_open
 local INCINFS, THAW, THAWN, CAPREF = E.incinfs, E.thawH, E.thawN, E.capref
 local GVAR, GCALL, GFINDALL = E.g_var, E.g_call, E.g_findall

prolog_engine.lua

@@ -76,7 +76,7 @@ local cell_top, var_top, trail_top, cap_top = 0, 0, 0, 0
 
 -- Kernel globals / error thrower, bound in install(). Declared up here
 -- because newvar's per-query capacity check needs them.
-local GLOBALS, ERRFN
+local GLOBALS, ERRFN, KEQ   -- KEQ: the kernel's `=`
 
 -- Per-query variable capacity (the kernel's shen.prolog-vector ceiling).
 -- shen.call-prolog allocates ONE absvector of (value shen.*prolog-memory*)
@@ -735,6 +735,40 @@ local function materialize(v)
 end
 M.materialize = materialize
 
+-- eqterm: Shen `=` of two arena terms, i.e. (= (materialize a)
+-- (materialize b)), without building either side. Identical terms are equal
+-- and a cons never equals a non-cons. Any other pair of distinct non-cons
+-- terms (atoms or vars) is decided by the kernel's `=` on their
+-- materializations, which are O(1) (an atom's value or a cached pvar), so
+-- opaque values such as tuples and vectors keep their structural `=`.
+-- Iterative over the cdr spine; car nesting recurses and is counted like
+-- import's (runtime.lua stack segments, #93).
+local eqterm
+eqterm = function(a, b)
+  while true do
+    a = lazyderef(a); b = lazyderef(b)
+    if a == b then return true end
+    if a >= CONS_BASE then
+      if b < CONS_BASE then return false end
+      local pa, pb = a - CONS_BASE, b - CONS_BASE
+      local ca, cb = cells[pa], cells[pb]
+      if ca ~= cb then
+        local same
+        local d = SD[1]
+        if d < SEG_FRAMES then SD[1] = d + 1; same = eqterm(ca, cb); SD[1] = d
+        else same = SEGRUN(eqterm, ca, cb) end
+        if not same then return false end
+      end
+      a = cells[pa + 1]; b = cells[pb + 1]
+    elseif b >= CONS_BASE then
+      return false
+    else
+      return KEQ(materialize(a), materialize(b))
+    end
+  end
+end
+M.eqterm = eqterm
+
 -- ---------------------------------------------------------------------------
 -- goal builtins (legacy 6-arg goal ABI becomes (args..., n, cont))
 -- ---------------------------------------------------------------------------
@@ -844,6 +878,7 @@ function M.install(P)
   M.P = P
   GLOBALS = P.GLOBALS
   ERRFN = P.ERR
+  KEQ = function(x, y) return P.F["="](x, y) end
   for _, mod in ipairs({ "prolog_compile", "typecheck_native" }) do
     local ok, m = pcall(require, mod)
     if ok then

scripts/run-tests.lua

@@ -34,6 +34,7 @@ local specs = {
   "test/checked_integer_spec.lua",
   "test/datatype_consume_spec.lua",
   "test/engine_spec.lua",
+  "test/eqterm_spec.lua",
   "test/error_robustness_spec.lua",
   "test/eval_order_spec.lua",
   "test/features_spec.lua",

test/eqterm_spec.lua

@@ -0,0 +1,151 @@
+-- test/eqterm_spec.lua — native arena equality for `=` guards.
+--
+--   luajit test/eqterm_spec.lua
+--
+-- E.eqterm(a, b) must agree with the kernel's (= (materialize a)
+-- (materialize b)) on every pair of terms, and the Prolog translator must
+-- route (= T1 T2) guards over engine terms through it, whether each side is
+-- spelled (shen.deref X B) or as a bare term variable (43.1's deref-hoisted
+-- by-hypothesis-h compares the bare, already-dereferenced X).
+package.path = (arg[0]:gsub("test/[^/]*$", "")) .. "?.lua;" .. package.path
+
+local shen = require("shen")
+shen.boot{ quiet = true }
+local R = require("runtime")
+local E = require("prolog_engine")
+local PC = require("prolog_compile")
+local F = shen.prims.F
+
+local npass, nfail = 0, 0
+local function check(cond, name)
+  if cond then npass = npass + 1
+  else
+    nfail = nfail + 1
+    io.write("FAIL: ", name, "\n")
+  end
+end
+
+local function kl(s) return R.read_all(s)[1] end
+
+-- ---------------------------------------------------------------------------
+-- eqterm agrees with `=` over materialized terms
+-- ---------------------------------------------------------------------------
+do
+  local q = E.query_begin()
+  local imp = function(x) return E.import(x, {}) end
+  local v1, v2 = E.newvar(), E.newvar()
+  local bound = E.newvar()
+  E.bind1(bound, imp(kl("(a b c)")), function() return true end)
+  local tup = F["@p"](1, 2)
+  local tup2 = F["@p"](1, 2)
+  local vec = F["@v"](1, F["vector"](0))
+  local vec2 = F["@v"](1, F["vector"](0))
+  local terms = {
+    imp(kl("a")), imp(kl("b")), imp(1), imp(1.0), imp(2), imp("s"), imp("t"),
+    imp(true), imp(false), imp(R.NIL),
+    imp(kl("(a b c)")), imp(kl("(a b c)")), imp(kl("(a b)")), imp(kl("(a b c d)")),
+    imp(kl("((a (b)) c)")), imp(kl("((a (b)) c)")), imp(kl("((a (c)) c)")),
+    imp(tup), imp(tup2), imp(F["@p"](1, 3)), imp(vec), imp(vec2),
+    v1, v2, bound,
+    E.cons(v1, 0), E.cons(v1, 0), E.cons(v2, 0),
+    E.cons(imp(kl("a")), v1),
+  }
+  local agree, total = 0, 0
+  for i = 1, #terms do
+    for j = 1, #terms do
+      total = total + 1
+      local want = F["="](E.materialize(terms[i]), E.materialize(terms[j]))
+      local got = E.eqterm(terms[i], terms[j])
+      if got == want then agree = agree + 1
+      else io.write("  eqterm mismatch at ", i, ",", j, ": got ", tostring(got),
+                    " want ", tostring(want), "\n") end
+    end
+  end
+  check(agree == total, "eqterm == (= (materialize a) (materialize b)) on all "
+        .. total .. " pairs")
+  check(E.eqterm(bound, imp(kl("(a b c)"))), "bound var derefs before compare")
+  check(E.eqterm(imp(tup), imp(tup2)), "distinct tuples compare structurally")
+
+  -- car nesting deeper than one Lua stack runs on stack segments (#93)
+  local function deep(n)
+    local t = imp(kl("x"))
+    for _ = 1, n do t = E.cons(t, 0) end
+    return t
+  end
+  local ok, r = pcall(E.eqterm, deep(200000), deep(200000))
+  check(ok and r == true, "200000-deep car nesting compares (" .. tostring(r) .. ")")
+  E.query_end(q)
+end
+
+-- ---------------------------------------------------------------------------
+-- the translator lifts (= T1 T2) guards over engine terms to EQTERM
+-- ---------------------------------------------------------------------------
+local function translated(src)
+  local fn, err, out = PC.translate_defun(kl(src))
+  return fn, err, out or ""
+end
+
+do
+  -- shape of 43.1 shen.by-hypothesis (both sides dereferenced)
+  local _, err, src = translated([[
+(defun spec.byh (V1 V2 V3 V4 V5 V6 V7)
+  (if (shen.unlocked? V5)
+      (let W1 (shen.lazyderef V3 V4)
+        (if (cons? W1)
+            (let W2 (hd W1)
+              (do (shen.incinfs)
+                  (when (= (shen.deref V1 V4) (shen.deref W2 V4)) V4 V5 V6
+                    (freeze (is! V2 W2 V4 V5 V6 V7)))))
+            false))
+      false))]])
+  check(src:find("EQTERM(", 1, true) ~= nil, "deref/deref guard -> EQTERM (" .. tostring(err) .. ")")
+  check(src:find("MAT(", 1, true) == nil, "deref/deref guard materializes nothing")
+end
+
+do
+  -- shape of the deref-hoisted shen.by-hypothesis-h: (= X (shen.deref Y B))
+  local _, err, src = translated([[
+(defun spec.byhh (V1 V2 V3 V4 V5 V6 V7)
+  (if (shen.unlocked? V5)
+      (let W1 (shen.lazyderef V3 V4)
+        (if (cons? W1)
+            (let W2 (hd W1)
+              (do (shen.incinfs)
+                  (when (= V1 (shen.deref W2 V4)) V4 V5 V6
+                    (freeze (is! V2 W2 V4 V5 V6 V7)))))
+            false))
+      false))]])
+  check(src:find("EQTERM(", 1, true) ~= nil, "bare/deref guard -> EQTERM (" .. tostring(err) .. ")")
+  check(src:find("MAT(", 1, true) == nil, "bare/deref guard materializes nothing")
+end
+
+do
+  -- any other guard still lifts to a KL defun over materialized arguments
+  local _, err, src = translated([[
+(defun spec.other (V1 V2 V3 V4 V5)
+  (if (shen.unlocked? V3)
+      (do (shen.incinfs)
+          (when (element? (shen.deref V1 V2) (cons a ())) V2 V3 V4 V5))
+      false))]])
+  check(src:find("MAT(", 1, true) ~= nil and src:find("EQTERM(", 1, true) == nil,
+        "non-= guard keeps the lifted path (" .. tostring(err) .. ")")
+end
+
+-- ---------------------------------------------------------------------------
+-- end to end: hypotheses are consulted through the native by-hypothesis
+-- ---------------------------------------------------------------------------
+do
+  local function ty(src, t)
+    local r = shen.typecheck(src, t)
+    return r ~= false and R.to_str(r) or false
+  end
+  check(ty("(/. X (+ X 1))", "(number --> number)") == "(number --> number)",
+        "lambda typechecks through hypotheses")
+  check(ty("(/. X Y (@p X Y))", "(A --> B --> (A * B))") ~= false,
+        "polymorphic pair typechecks")
+  check(ty("(/. X (+ X \"a\"))", "(number --> number)") == false,
+        "ill-typed body still fails")
+end
+
+io.write(string.format("eqterm_spec: %d pass, %d fail\n", npass, nfail))
+os.exit(nfail == 0 and 0 or 1)