Multi-prompted Shift/Reset in metalua

shiftreset.lua

-- Multi-prompted shift/reset in metalua
-- after Roshan James & Amr Sabry @ http://parametricity.net/dropbox/yield.subc.pdf
-- indepth comments and explanation are to be found in the paper.
-- The addition of multi-prompt is mine (Arnaud Spiwack) and documented in the
-- only function I had to change significantly from the single-prompted one
-- (namely runco). The rest is just carrying the prompt around.
-- The implementation is ok as long as continuations are not used twice.
-- To be work properly in all generality, the coroutine.clone primitive
-- must be implemented. There exists a patch against Lua 5.1 to that effect.
-- Hopefully it'll make it into the distribution soon.

-- Should be a .mlua file, but I .lua-ed it to please github's syntax highlighter

-{extension 'match'}

function clone (x) return x end --should be coroutine.clone
function snd(x,y) return y end -- to access the relevant return value of resume

-- run & case allow to distinguish between values that were yielded and
-- those which have been returned.
function case (x,res,ret) match x with
  | `Susp {y, cont} -> return res (y, cont)
  | `Return {x} -> return ret (x)
  end
end

function runco (prompt,co,i)
  match snd (coroutine.resume(co,i)) with
  | `Return {x} -> return `Return {x}
  | `Yielded {p,o} ->
     local function res (i2)
	local co2 = clone(co)
	return runco(prompt,co2,i2)
     end
     if p == prompt then
	return `Susp { o , res }
     else
	-- this is where I handle multi-prompt. If a runco
        -- receives a yield that it's not supposed to control, it simply
	-- re-yields it and when it gets restarted it immediatly restarts
        -- its coroutine, propagating the argument.
        -- It may take a little time to understand why this restores the
        -- state we expect. The intuition is that the runco are stacked
        -- and control a single coroutine. So yielding here turns the control
        -- to the previous runco which will create (if it matches the prompt)
        -- a handle to restart here.
	return res (coroutine.yield(`Yielded {p,o}))
     end
  end
end

function run (prompt,b)
  local co = coroutine.create(function () return `Return {b()} end)
  return runco(prompt,co)
end

prompt_count = 0
function new_prompt ()
  prompt_count = prompt_count+1
  return prompt_count
end

-- shift and reset
function shift (prompt,f)
   return coroutine.yield(`Yielded{prompt,f})
end

function reset (prompt,b)
  return rereset (prompt,run (prompt,b))
end

function rereset (prompt,b)
  return case (b,
	       function (f,k) 
		  return reset (prompt,
				function () return f (function (i) return rereset(prompt,k(i)) end) end)
	       end,
	       function (x) return x end
  )
end

-- test: counters
function new_counter () return new_prompt() end

function incr (c)
   return shift (c,function (k)
		    return function (i) return k({})(i+1) end
		 end)
end

function get (c)
   return shift (c,function (k)
		    return function (i) return k(i)(i) end
		 end)
end

function counter(c,body)
   return reset (c,function () local r = body (); return function (i) return r end end) (0)
end

c1 = new_counter()

c2 = new_counter()

print (counter(c1, function () return get(c1) end)) -- 0
print (counter(c1, function () incr(c1) return get(c1) end)) -- 1
print (counter(c1, function () incr(c1) incr(c1) return get(c1) end)) -- 2

print (counter(c1, function () return counter(c2, function ()
						     incr(c1) incr(c1) incr(c1)
						     incr(c2) incr(c2)
						     return get(c1)+get(c2)
						  end) end)) -- 5