fix: run deep non-tail recursion on fresh coroutine stacks (#93; builds on #90)

CHANGELOG.md

@@ -39,6 +39,19 @@ Notable changes to shen-lua. Release notes for earlier versions are under
 
 ### Fixed
 
+- Deep non-tail recursion no longer overflows the Lua stack, in user code
+  or in the kernel's own tree walkers. `(depth 12000)` for a plain
+  `(+ 1 (depth (- N 1)))` recursion failed with `stack overflow` (it reached
+  ~8000), as did typechecking a 3000-deep `[s N]` term or a 10000-element
+  `[list number]` (`shen.curry`) and reading a 4000-element list literal
+  (`shen.find-types`, `shen.process-applications`). LuaJIT's stack is a
+  fixed ~65500 slots per coroutine (`LJ_STACK_MAX`, a compile-time
+  constant), so a compiled function that makes a non-tail call now counts
+  itself while it is on the stack, and every 160 of them the next call runs
+  on a fresh coroutine stack. Tail calls are not counted, so constant-stack
+  loops are unchanged. Depth is bounded by memory (`(depth 100000)` works);
+  past 10000 nested segments, shared with the Prolog engine's, a call raises
+  the trappable `Lua stack depth limit exceeded` (pyrex41/shen-lua#93).
 - A warm fasl hit no longer runs a nested `load` twice. When a loaded file
   loaded another file, from a top-level `(load ...)` form or from a function
   that loads, the nested file's output printed twice on every cached run.

bin/shen

@@ -83,7 +83,7 @@ if #actions == 0 then
     while not stin.eof do
       F["shen.initialise_environment"]()
       F["shen.prompt"]()
-      local okk, err = pcall(F["shen.read-evaluate-print"])
+      local okk, err = R.seg.pcall(F["shen.read-evaluate-print"])
       if not okk and not stin.eof then
         local e = getmetatable(err) == R.Excn and err or R.mkexcn(tostring(err))
         pcall(F["shen.toplevel-display-exception"], e)

boot.lua

@@ -3,6 +3,8 @@
 -- ShenOSKernel 43.1 modules are pure defuns; initialise() below runs
 -- (shen.initialise), the extension initialisers and (stlib.initialise).
 local R = require("runtime")
+-- restores the stack-segment count on a caught error (runtime.lua, #93)
+local pcall = R.seg.pcall
 local C = require("compiler")
 local P = require("prims")
 

lua_interop.lua

@@ -629,7 +629,8 @@ function M.call(name, ...)
   local n = select("#", ...)
   local a = { ... }
   for i = 1, n do a[i] = to_shen(a[i]) end
-  return to_lua(P.APP(fn, unpack(a, 1, n)))
+  -- R.seg.call: a caught error leaves the stack-segment count intact (#93)
+  return to_lua(R.seg.call(P.APP, fn, unpack(a, 1, n)))
 end
 
 function M.fn(name)

prims.lua

@@ -7,6 +7,11 @@ local cons, is_cons, NIL = R.cons, R.is_cons, R.NIL
 local Symbol, intern, is_symbol = R.Symbol, R.intern, R.is_symbol
 local Excn, mkexcn = R.Excn, R.mkexcn
 local Vmt = R.Vmt        -- absvector metatable (pure-array layout; shared with runtime.lua)
+-- counted non-tail calls (runtime.lua stack segments, #93), written inline
+-- where a native calls back into Shen code (a shared trampoline would look
+-- like recursion to the trace recorder)
+local SD, SEGRUN, SEG_FRAMES = R.seg.SD, R.seg.run, R.seg.FRAMES
+local pcall = R.seg.pcall  -- restores the segment count on a caught error
 
 local P = {}
 
@@ -1129,7 +1134,12 @@ function P.install_native_stdlib()
     while true do
       if lst == NIL then return reverse_help(acc, NIL) end
       if is_cons(lst) then
-        acc = cons(APP(f, lst[1]), acc); lst = lst[2]
+        -- a counted call: a walker that recurses through map nests here
+        local v
+        local d = SD[1]
+        if d < SEG_FRAMES then SD[1] = d + 1; v = APP(f, lst[1]); SD[1] = d
+        else v = SEGRUN(APP, f, lst[1]) end
+        acc = cons(v, acc); lst = lst[2]
       else
         return orig_maph(f, lst, acc)
       end
@@ -1812,7 +1822,10 @@ function P.install_native_stdlib()
         if last then return head else return NIL end
       end
       if not is_cons(lst) then return orig_mapcan(f, orig) end
-      local chunk = APP(f, lst[1])
+      local chunk
+      local d = SD[1]   -- a counted call (see map_h)
+      if d < SEG_FRAMES then SD[1] = d + 1; chunk = APP(f, lst[1]); SD[1] = d
+      else chunk = SEGRUN(APP, f, lst[1]) end
       if chunk ~= NIL then
         -- append copies the left spine; copy each chunk so we don't mutate
         -- a list the mapper returned that might be shared. A non-list chunk
@@ -2630,13 +2643,19 @@ local ENV = {
   ADSET = function(v, i, x) v[i + 2] = x; return v end,
   VMT = Vmt,
   -- allow compiled code to reach a few Lua builtins safely
-  pcall = pcall, select = select, error = error,
+  -- trap-error's pcall restores the stack-segment count (runtime.lua, #93)
+  pcall = R.seg.pcall, select = select, error = error,
   setmetatable = setmetatable, getmetatable = getmetatable,
   type = type,
   math = math, string = string, table = table,
 }
 P.ENV = ENV
 
+-- counted non-tail calls (runtime.lua stack segments, #93): the count slot,
+-- the segment runner and the trampoline for code outside a defun
+ENV.SD, ENV.SEG, ENV.SEGSELF, ENV.NTV = R.seg.SD, R.seg.run, R.seg.runself, R.seg.NTV
+R.seg.set_error(ERR)
+
 -- CALL0..CALLn: late-bound fixed-arity calls to a function NAMED at compile
 -- time whose arity was not yet known (forward references; see unknown_call
 -- in compiler.lua). Generated from one template so each is a plain

prolog_engine.lua

@@ -41,6 +41,7 @@ if not ok_ffi then
   return { install = function() end }
 end
 local R = require("runtime")
+local SD, SEGRUN, SEG_FRAMES = R.seg.SD, R.seg.run, R.seg.FRAMES
 
 local loadstring = loadstring or load   -- 5.2+ shim (module is FFI-gated anyway)
 
@@ -242,10 +243,6 @@ local SEG_FRAMES = 96        -- counted frames per segment: a typecheck
 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
@@ -257,17 +254,13 @@ local function thaw_segment(cont)
   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
+  -- the runtime's segment runner (runtime.lua, #93): its nesting cap covers
+  -- both kinds of segment, which share one C stack
+  local ok, r = pcall(SEGRUN, thawH, cont)
   -- 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]
+  if not ok then error(r, 0) end
+  return r
 end
 
 -- thawH for a non-tail site: counts the frame, or opens a segment. The
@@ -641,7 +634,16 @@ local function import(x, varmap)
     end
     local tail = import(x, varmap)
     for i = sn, 1, -1 do
-      tail = mkcons(import(spine[i], varmap), tail)
+      local e = spine[i]
+      if getmt(e) == Cons then
+        -- car nesting recurses: count it (runtime.lua stack segments, #93)
+        local d = SD[1]
+        if d < SEG_FRAMES then SD[1] = d + 1; e = import(e, varmap); SD[1] = d
+        else e = SEGRUN(import, e, varmap) end
+      else
+        e = import(e, varmap)
+      end
+      tail = mkcons(e, tail)
     end
     return tail
   elseif getmt(x) == Vmt and x[2] == shen_pvar then

repl.lua

@@ -324,7 +324,8 @@ function M.run(opts)
     end
     if buf:match("%S") then
       ed.remember(buf)
-      local ok, r = xpcall(function() return eval_input(buf) end, handler)
+      -- the runtime's xpcall restores the stack-segment count (#93)
+      local ok, r = require("runtime").seg.xpcall(function() return eval_input(buf) end, handler)
       if not ok then
         -- r is the handler capture, unless something threw inside the
         -- handler itself (then r is whatever Lua gives us).

runtime.lua

@@ -292,4 +292,102 @@ M.to_str = to_str
 M.shortest_float = shortest_float
 M.num_to_str = num_to_str
 
+----------------------------------------------------------------------
+-- Lua stack segments (pyrex41/shen-lua#93)
+----------------------------------------------------------------------
+-- LuaJIT's Lua stack is a fixed ~65500 slots per coroutine (LJ_STACK_MAX,
+-- a compile-time constant), so non-tail recursion in compiled KL -- the
+-- kernel's own tree walkers (shen.curry, shen.find-types,
+-- shen.process-applications) and user code alike -- overflowed after about
+-- 12k frames, where shen-go grows its stack. Every coroutine has its own
+-- stack, so deep recursion continues on a fresh one (a segment), as the
+-- Prolog engine already does for deep proofs (#85).
+--
+-- SD[1] counts the compiled functions on the current stack that are waiting
+-- on a non-tail call (compiler.lua NT_SRC: such a function counts itself on
+-- entry and restores the count before it returns or tail-calls). Once
+-- SEG_FRAMES of them are open, the next one runs on a new coroutine instead
+-- (M.seg.run). Tail calls are never counted, so a loop that runs in constant
+-- stack never opens a segment. NTV is the trampoline form for code outside
+-- a defun; natives that call back into Shen code count inline.
+--
+-- An error unwinds past the restores, so the catch points restore the count
+-- by assignment: M.seg.pcall / .xpcall (the compiled trap-error and the
+-- runtime's own pcalls use them), M.seg.call (host entry points, so an error
+-- a host catches does not leak the count) and the segment runner itself.
+do
+  local co_create, co_resume, co_status, co_yield =
+    coroutine.create, coroutine.resume, coroutine.status, coroutine.yield
+  local unpack = unpack or table.unpack
+  local SEG_FRAMES = 160      -- counted calls per segment (C.SEG_FRAMES)
+  local SEG_MAX = 10000       -- nested segments: each costs ~200 bytes of C stack
+  local SD = { 0 }            -- counted calls open on the current stack
+  local nseg, total = 0, 0
+  local ERR = function(msg) error(msg, 0) end
+  local function pack(...) return { n = select("#", ...), ... } end
+  local S = { SD = SD, FRAMES = SEG_FRAMES }
+
+  -- Run f(...) 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.
+  function S.run(f, ...)
+    if nseg >= SEG_MAX then
+      ERR("Lua stack depth limit exceeded (" .. SEG_MAX .. " stack segments)")
+    end
+    local saved, saved_n = SD[1], nseg
+    nseg = nseg + 1; SD[1] = 0
+    total = total + 1
+    local co = co_create(f)
+    local res = pack(co_resume(co, ...))
+    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
+    SD[1], nseg = saved, saved_n
+    if not res[1] then error(res[2], 0) end
+    return res[2]
+  end
+  local run = S.run
+  -- S.run for the function that called it (a lambda has no name to pass)
+  local getinfo = debug and debug.getinfo
+  function S.runself(...)
+    return run(getinfo(2, "f").func, ...)
+  end
+
+  -- the trampoline outside a defun (vararg: a FUNCV entry never starts a
+  -- root trace of its own)
+  function S.NTV(f, ...)
+    local d = SD[1]
+    if d < SEG_FRAMES then SD[1] = d + 1; local r = f(...); SD[1] = d; return r end
+    return run(f, ...)
+  end
+
+  local function restore(d, ok, ...)
+    if not ok then SD[1] = d end
+    return ok, ...
+  end
+  local function rethrow(d, ok, ...)
+    if not ok then SD[1] = d; error((...), 0) end
+    return ...
+  end
+  function S.pcall(f, ...) return restore(SD[1], pcall(f, ...)) end
+  function S.xpcall(f, h) return restore(SD[1], xpcall(f, h)) end
+  -- f(...) for a host entry point
+  function S.call(f, ...) return rethrow(SD[1], pcall(f, ...)) end
+
+  -- segments ever opened, nested now, counted calls open on this stack
+  function S.stats() return total, nseg, SD[1] end
+  -- set the nesting cap (tests); returns the previous value
+  function S.limit(max)
+    local m0 = SEG_MAX
+    SEG_MAX = max or m0
+    return m0
+  end
+  function S.set_error(fn) ERR = fn end
+  M.seg = S
+end
+
 return M

scripts/run-tests.lua

@@ -49,6 +49,7 @@ local specs = {
   "test/reader_native_spec.lua",
   "test/reader_spec.lua",
   "test/repl_spec.lua",
+  "test/stack_segments_spec.lua",
   "test/stdlib_spec.lua",
   "test/tailcall_spec.lua",
   "test/typecheck_api_spec.lua",

shen.lua

@@ -12,6 +12,8 @@
 -- machinery is reimplemented.
 
 local R = require("runtime")
+-- restores the stack-segment count on a caught error (runtime.lua, #93)
+local pcall = R.seg.pcall
 local P = require("boot")   -- boot.lua returns the prims module P, fully wired
 local checked_integer = require("checked_integer")
 
@@ -79,8 +81,7 @@ end
 -- with the kernel's own reader (read-from-string -> list of forms) and (eval)
 -- each form through the real macroexpand/shen->kl pipeline. Returns the value
 -- of the last form (a Shen value: number/string/boolean/symbol/cons/...).
-function shen.eval(src)
-  ensure_boot()
+local function eval_forms(src)
   local forms = P.F["read-from-string"](src)
   local last
   while R.is_cons(forms) do
@@ -90,6 +91,11 @@ function shen.eval(src)
   return last
 end
 
+function shen.eval(src)
+  ensure_boot()
+  return R.seg.call(eval_forms, src)
+end
+
 -- ---- call / fn ------------------------------------------------------------
 -- Call the Shen function `name` with Lua values as arguments. Arity is
 -- handled by the existing P.APP machinery: fewer args than the function's
@@ -99,7 +105,7 @@ function shen.call(name, ...)
   ensure_boot()
   local fn = P.F[name]
   if fn == nil then error("shen.call: undefined function: " .. tostring(name), 2) end
-  return P.APP(fn, ...)
+  return R.seg.call(P.APP, fn, ...)
 end
 
 -- A plain Lua callable for the Shen function `name`. The F-table lookup is