fix(prolog): run deep proofs on fresh coroutine stacks (#85)

CHANGELOG.md

@@ -39,6 +39,17 @@ Notable changes to shen-lua. Release notes for earlier versions are under
 
 ### Fixed
 
+- The native Prolog engine no longer overflows the Lua stack on deep
+  proofs: a typecheck died with `stack overflow` at ~375 nested rule
+  applications (`(shen.typecheck (mk 400) peano)`, an 800-element
+  `[list number]`), well short of the legacy engine's ~1000. The engine runs
+  in continuation-passing style on the Lua stack, and LuaJIT's stack is a
+  fixed ~65500 slots per coroutine, so a deep proof now continues on a fresh
+  coroutine stack every 96 counted non-tail continuation calls. Depth is
+  bounded by memory; past 1000 nested segments a query raises the trappable
+  `Prolog depth limit exceeded`. Typechecks of 1000-20000-deep terms now stop
+  at the kernel's own recursive preprocessing (`shen.curry`, ~2000 nested
+  `[s N]` levels or ~4600 list elements) instead (pyrex41/shen-lua#85).
 - The kernel bytecode cache and its stdlib-signature sidecar default to
   `${XDG_CACHE_HOME:-$HOME/.cache}/shen-lua/kernel-<build>.bin` (directory
   created on first write; an unwritable location just skips caching). They

prolog_compile.lua

@@ -43,6 +43,33 @@ local C = require("compiler")
 local unpack = table.unpack or unpack
 local loadstring = loadstring or load
 
+-- A value in return position: `THAWN(x)` (a counted, non-tail thaw) as the
+-- whole expression becomes `THAW(x)`, a proper tail call that neither grows
+-- the stack nor counts toward a segment (prolog_engine.lua, #85).
+local function tail_return(val)
+  if val:sub(1, 6) ~= "THAWN(" then return val end
+  local depth, i, n = 0, 6, #val
+  while i <= n do
+    local c = val:sub(i, i)
+    if c == '"' then
+      i = i + 1
+      while i <= n and val:sub(i, i) ~= '"' do
+        if val:sub(i, i) == "\\" then i = i + 1 end
+        i = i + 1
+      end
+    elseif c == "(" then depth = depth + 1
+    elseif c == ")" then
+      depth = depth - 1
+      if depth == 0 then
+        if i == n then return "THAW(" .. val:sub(7) end
+        return val
+      end
+    end
+    i = i + 1
+  end
+  return val
+end
+
 local M = {}
 
 local E, P, F            -- bound in install()
@@ -458,7 +485,9 @@ compile_value = function(ctx, e, env, out, d)
         return ctx:callstatic(b, {}, env, out, d)
       end
     end
-    return "THAW(" .. compile_value(ctx, x, env, out, d) .. ")"
+    -- THAWN counts the frame toward the engine's stack segments (#85);
+    -- tail_return turns it back into a plain THAW in return position
+    return "THAWN(" .. compile_value(ctx, x, env, out, d) .. ")"
   end
 
   -- (shen.gc B E): LIFO-pop the var allocated by the paired shen.newpv when
@@ -563,7 +592,7 @@ compile_value = function(ctx, e, env, out, d)
         local k = compile_cont(ctx, a[an], env, out, d)
         local r = newlocal(ctx)
         out[#out + 1] = indent(d) .. "local " .. r .. " = false"
-        out[#out + 1] = indent(d) .. "if " .. tl .. " then " .. r .. " = THAW(" .. k .. ") end"
+        out[#out + 1] = indent(d) .. "if " .. tl .. " then " .. r .. " = THAWN(" .. k .. ") end"
         return r
       elseif nm == "is" or nm == "is!" or nm == "bind" then
         local x = compile_term(ctx, a[1], env, out, d)
@@ -701,7 +730,7 @@ function CtxMT.lift_static(ctx, name, params, body, env)
   local fval = compile_value(ctx, body, penv, fbuf, 1)
   local def = { fname .. " = function(" .. table.concat(sig, ", ") .. ")" }
   for _, l in ipairs(fbuf) do def[#def + 1] = l end
-  def[#def + 1] = "  return " .. fval
+  def[#def + 1] = "  return " .. tail_return(fval)
   def[#def + 1] = "end"
   ctx.buf[#ctx.buf + 1] = table.concat(def, "\n")
   return st
@@ -763,7 +792,7 @@ function CtxMT.lift_cont(ctx, body, env, out, d)
   local def = { fname .. " = function(base, h)" }
   for _, l in ipairs(decls) do def[#def + 1] = "  " .. l end
   for _, l in ipairs(fbuf) do def[#def + 1] = l end
-  def[#def + 1] = "  return " .. fval
+  def[#def + 1] = "  return " .. tail_return(fval)
   def[#def + 1] = "end"
   ctx.buf[#ctx.buf + 1] = table.concat(def, "\n")
 
@@ -803,7 +832,7 @@ 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 CUT, UNLOCK, LOCKOPEN = E.cut, E.unlock, E.lock_is_open
-local INCINFS, THAW, CAPREF = E.incinfs, E.thawH, E.capref
+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
 local NEWCONT0, NEWCONT1, NEWCONT2, NEWCONT3, NEWCONT4 =
   E.newcont0, E.newcont1, E.newcont2, E.newcont3, E.newcont4
@@ -865,7 +894,7 @@ local function translate_core(name, allparams, body)
   for _, def in ipairs(ctx.buf) do src[#src + 1] = def end
   src[#src + 1] = "return function(" .. table.concat(sig, ", ") .. ")"
   for _, l in ipairs(mainbuf) do src[#src + 1] = l end
-  src[#src + 1] = "  return " .. val
+  src[#src + 1] = "  return " .. tail_return(val)
   src[#src + 1] = "end"
   local source = table.concat(src, "\n")
 

prolog_engine.lua

@@ -217,6 +217,72 @@ local function thawH(h)
 end
 M.thawH = thawH
 
+-- Stack segments (pyrex41/shen-lua#85). The engine is CPS on the Lua stack:
+-- a successful proof returns only once the whole query has succeeded, so its
+-- depth is the Lua stack depth. LuaJIT's stack is a fixed ~65500 slots per
+-- coroutine, which capped a typecheck at ~375 nested rule applications. But
+-- every coroutine has its own stack, so a deep proof continues in a fresh
+-- coroutine (a segment) and depth becomes bounded by memory instead.
+--
+-- The trigger counts the engine's own NON-tail continuation calls -- the
+-- thaws in unify, bind1 and cut, which must stay non-tail because they undo
+-- bindings / set the lock when the continuation fails. Every rule
+-- application on a proof path passes through at least one of them, and the
+-- generated predicate frames between two of them are bounded by the clause
+-- shapes. Tail-called thaws never count, so a loop that runs in constant
+-- stack never opens a segment.
+local SEG_FRAMES = 96        -- counted frames per segment: a typecheck
+                             -- overflowed one stack after ~750 of them
+local SEG_MAX = 1000         -- nested segments per query before giving up
+local seg_depth, seg_count = 0, 0
+local seg_total = 0          -- segments ever opened (M.segments, diagnostics)
+local co_create, co_resume, co_status, co_yield =
+  coroutine.create, coroutine.resume, coroutine.status, coroutine.yield
+local unpack = unpack or table.unpack
+local function pack(...) return { n = select("#", ...), ... } end
+
+-- Thaw `cont` on a fresh coroutine stack and return its value. Errors are
+-- rethrown unchanged (Shen exception objects included); a yield from inside
+-- (an embedding host's) is passed through to the enclosing coroutine.
+local function thaw_segment(cont)
+  if seg_count >= SEG_MAX then
+    ERRFN("Prolog depth limit exceeded (" .. SEG_MAX .. " stack segments)")
+  end
+  local saved, saved_count = seg_depth, seg_count
+  seg_count = seg_count + 1; seg_depth = 0
+  seg_total = seg_total + 1
+  local co = co_create(thawH)
+  local res = pack(co_resume(co, cont))
+  while res[1] and co_status(co) ~= "dead" do
+    local up = pack(pcall(co_yield, unpack(res, 2, res.n)))
+    if not up[1] then res = up; break end   -- no coroutine to yield to
+    res = pack(co_resume(co, unpack(up, 2, up.n)))
+  end
+  -- assign, not decrement: segments abandoned while suspended never return
+  seg_count, seg_depth = saved_count, saved
+  if not res[1] then error(res[2], 0) end
+  return res[2]
+end
+
+-- thawH for a non-tail site: counts the frame, or opens a segment. The
+-- translated predicates call it (THAWN) wherever a thaw is not their return
+-- value. unify, bind1 and cut write the same lines out inline instead, so
+-- the counting costs them no extra frame:
+--   local d = seg_depth
+--   local r
+--   if d < SEG_FRAMES then seg_depth = d + 1; r = thawH(cont); seg_depth = d
+--   else r = thaw_segment(cont) end
+local function thawN(h)
+  local d = seg_depth
+  if d >= SEG_FRAMES then return thaw_segment(h) end
+  seg_depth = d + 1
+  local r
+  if type(h) == "number" then r = contFn[h](contBase[h], h) else r = h() end
+  seg_depth = d
+  return r
+end
+M.thawN = thawN
+
 local function ckroom(n)
   if cap_top + n > kcap then capbuf, kcap = grown(capbuf, kcap, cap_top + n) end
 end
@@ -463,7 +529,10 @@ local function unify_core(rootA, rootB, cont, oc)
   end
   if ok then
     trail_top = tn
-    local r = thawH(cont)
+    local d = seg_depth
+    local r
+    if d < SEG_FRAMES then seg_depth = d + 1; r = thawH(cont); seg_depth = d
+    else r = thaw_segment(cont) end
     if r == false then
       for i = tn - 1, tmark, -1 do vbind[trail[i]] = -1 end
       trail_top = tmark
@@ -487,7 +556,10 @@ function M.bind1(varv, val, cont)
   if trail_top >= tcap then trail, tcap = grown(trail, tcap, trail_top + 1) end
   trail[trail_top] = idx
   trail_top = trail_top + 1
-  local r = thawH(cont)
+  local d = seg_depth
+  local r
+  if d < SEG_FRAMES then seg_depth = d + 1; r = thawH(cont); seg_depth = d
+  else r = thaw_segment(cont) end
   if r == false then
     vbind[idx] = -1
     trail_top = trail_top - 1
@@ -503,7 +575,10 @@ function M.lock_is_open() return lock_open end
 
 -- shen.cut: thaw; on failure with an open lock, lock at depth n.
 function M.cut(n, cont)
-  local r = thawH(cont)
+  local d = seg_depth
+  local r
+  if d < SEG_FRAMES then seg_depth = d + 1; r = thawH(cont); seg_depth = d
+  else r = thaw_segment(cont) end
   if r == false and lock_open then
     lock_open = false
     lock_depth = n
@@ -667,7 +742,7 @@ function M.query_begin()
   local q = {
     cell_top, var_top, trail_top, cap_top, ch_top,
     lock_open, lock_depth, opq_top,
-    var_base, var_limit,
+    var_base, var_limit, seg_depth, seg_count,
   }
   lock_open, lock_depth = true, 0
   -- Fresh bindings vector for this query, exactly as shen.call-prolog does:
@@ -699,6 +774,9 @@ function M.query_end(q)
   end
   opq_top = q[8]
   var_base, var_limit = q[9], q[10]
+  -- an error unwinds past the decrements in unify/bind1/cut; a nested query
+  -- shares its caller's stack, so it continues the caller's count
+  seg_depth, seg_count = q[11], q[12]
   -- the import memo references cells that are no longer valid
   if next(import_memo) ~= nil then
     for k in pairs(import_memo) do import_memo[k] = nil end
@@ -708,6 +786,7 @@ end
 -- full reset (tests / benchmarks only)
 function M.reset_all()
   cell_top, var_top, trail_top, cap_top, ch_top = 0, 0, 0, 0, 0
+  seg_depth, seg_count = 0, 0
   lock_open, lock_depth = true, 0
   var_base, var_limit = 0, math.huge
   infs = 0
@@ -721,6 +800,15 @@ function M.tops()
   return cell_top, var_top, trail_top, cap_top, ch_top
 end
 
+-- stack segments: ever opened, currently nested, counted frames in this one
+function M.segments() return seg_total, seg_count, seg_depth end
+-- set the segment size / nesting cap (tests); returns the previous values
+function M.segment_limits(frames, max)
+  local f0, m0 = SEG_FRAMES, SEG_MAX
+  SEG_FRAMES, SEG_MAX = frames or f0, max or m0
+  return f0, m0
+end
+
 -- ---------------------------------------------------------------------------
 -- predicate dispatch table (filled by prolog_compile.lua / typecheck_native)
 -- ---------------------------------------------------------------------------

test/prolog_semantics_spec.lua

@@ -165,6 +165,96 @@ check("var? under findall",
 check("var? failure backtracks into the next clause",
       q("(prolog? (is X 1) (tvarness X R) (return R))"), "bound")
 
+-- ---------------------------------------------------------------------------
+-- Proof depth (pyrex41/shen-lua#85). The native engine is CPS on the Lua
+-- stack, and LuaJIT's stack is a fixed ~65500 slots: a typecheck overflowed
+-- at ~375 nested peano rules (legacy engine ~1000; shen-go has no limit). A
+-- deep proof now continues on fresh coroutine stacks ("segments"). These
+-- check the depth, that backtracking and cut undo bindings across segment
+-- boundaries, and that errors cross them unchanged.
+-- ---------------------------------------------------------------------------
+do
+  local E = require("prolog_engine")
+  local seg0 = E.segments()
+  shen.eval([==[
+(datatype dpeano
+  ____________
+  z : dpeano;
+
+  N : dpeano;
+  ____________
+  [s N] : dpeano;)
+(define dmk N Acc -> (if (= N 0) Acc (dmk (- N 1) [cons s [cons Acc []]])))
+(define dlist N Acc -> (if (= N 0) Acc (dlist (- N 1) [cons N Acc])))
+(defprolog dchain
+  0 R <-- (is R bottom) (when false);
+  N R <-- (when (> N 0)) (is M (- N 1)) (dchain M R);
+  _ R <-- (is R top);)
+(defprolog dcut
+  0 <-- ! (when false);
+  0 <--;
+  N <-- (when (> N 0)) (is M (- N 1)) ! (dcut M);)
+(defprolog derr
+  0 <-- (when (= (/ 1 0) 0));
+  N <-- (when (> N 0)) (is M (- N 1)) (derr M);)]==])
+  check("#85: typecheck of a 1500-deep peano term",
+        R.to_str(shen.eval("(shen.typecheck (dmk 1500 z) dpeano)")), "dpeano")
+  check("#85: typecheck of a 3000-element (list number)",
+        R.to_str(shen.eval("(shen.typecheck (dlist 3000 []) [list number])")), "(list number)")
+  check("#85: the deep typechecks ran on more than one stack segment",
+        E.segments() > seg0, true)
+  local seg1 = E.segments()
+  check("#85: a shallow typecheck opens no segment",
+        R.to_str(shen.eval("(shen.typecheck (dmk 3 z) dpeano)")) == "dpeano"
+        and E.segments() == seg1, true)
+  -- the leaf binds R, fails, and every level's next clause is tried on the
+  -- way back up: R must be unbound again before (is R top)
+  check("#85: backtracking out of a 3000-deep failing branch undoes its bindings",
+        q("(prolog? (dchain 3000 R) (return R))"), "top")
+  check("#85: cut 3000 levels deep still prunes (no second clause for 0)",
+        q("(prolog? (dcut 3000))"), "failed")
+  local shallow = shen.eval("(trap-error (prolog? (derr 1)) (/. E (error-to-string E)))")
+  local deep = shen.eval("(trap-error (prolog? (derr 3000)) (/. E (error-to-string E)))")
+  check("#85: an error 3000 levels deep crosses segments unchanged",
+        type(deep) == "string" and deep == shallow, true)
+  check("#85: the engine is usable after an error inside a segment",
+        q("(prolog? (dchain 50 R) (return R))"), "top")
+  local _, nested = E.segments()
+  check("#85: no segment is left open", nested, 0)
+  -- a host coroutine.yield from inside a segment reaches the host's
+  -- coroutine and the proof resumes where it left off
+  do
+    local P = shen.prims
+    P.F["dyield"] = function() return coroutine.yield("paused") == "go" end
+    P.FA[P.F["dyield"]] = 0
+    shen.eval([==[(defprolog dyc
+  0 R <-- (when (dyield)) (is R bottom);
+  N R <-- (when (> N 0)) (is M (- N 1)) (dyc M R);)]==])
+    local co = coroutine.create(function()
+      return shen.eval("(prolog? (dyc 3000 R) (return R))")
+    end)
+    local ok1, v1 = coroutine.resume(co)
+    local ok2, v2 = coroutine.resume(co, "go")
+    check("#85: a yield from 3000 levels deep reaches the host and resumes",
+          tostring(ok1) .. " " .. tostring(v1) .. " " .. tostring(ok2) .. " "
+          .. R.to_str(v2), "true paused true bottom")
+    -- with no host coroutine the yield is an error, and abandoning the
+    -- suspended segments must not leak the nesting count
+    local okn = pcall(shen.eval, "(prolog? (dyc 3000 R) (return R))")
+    local _, nested2 = E.segments()
+    check("#85: a yield with no host coroutine raises and leaves no segment open",
+          tostring(okn) .. " " .. nested2, "false 0")
+  end
+  -- past the cap: a trappable Shen error, not a crash
+  local f0, m0 = E.segment_limits(8, 2)
+  local capped = shen.eval("(trap-error (prolog? (dchain 3000 R) (return R)) (/. E (error-to-string E)))")
+  E.segment_limits(f0, m0)
+  check("#85: past the segment cap the query raises a Shen error",
+        type(capped) == "string" and capped:find("Prolog depth limit exceeded", 1, true) ~= nil, true)
+  check("#85: and the next query runs normally",
+        q("(prolog? (dchain 3000 R) (return R))"), "top")
+end
+
 -- ---------------------------------------------------------------------------
 -- D-3: per-query variable capacity
 --