-{ extension 'match' }
-{ extension 'log' }

require 'metalua.walk'

----------------------------------------------------------------------
-- Back-end:
----------------------------------------------------------------------

-- Parse additional elements in a loop
loop_element = gg.multisequence{
   { 'while',  mlp.expr, builder = |x| `Until{ `Op{ 'not', x[1] } } },
   { 'until',  mlp.expr, builder = |x| `Until{ x[1] } },
   { 'if',     mlp.expr, builder = |x| `If{ x[1] } },
   { 'unless', mlp.expr, builder = |x| `If{ `Op{ 'not', x[1] } } },
   { 'for',    mlp.for_header, builder = |x| x[1] } }

-- Recompose the loop
function xloop_builder(x)
   local first, elements, body = unpack(x)

   -------------------------------------------------------------------
   -- If it's a regular loop, don't bloat the code
   -------------------------------------------------------------------
   if not next(elements) then
      table.insert(first, body)
      return first
   end

   -------------------------------------------------------------------
   -- There's no reason to treat the first element in a special way
   -------------------------------------------------------------------
   table.insert(elements, 1, first)

   -------------------------------------------------------------------
   -- if a header or a break must be able to exit the loops, ti will
   -- set exit_label and use it (a regular break wouldn't be enough,
   -- as it couldn't escape several nested loops.)
   -------------------------------------------------------------------
   local exit_label
   local function exit()
      if not exit_label then exit_label = mlp.gensym 'break' [1] end
      return `Goto{ exit_label }
   end

   -------------------------------------------------------------------
   -- Compile all headers elements, from last to first
   -------------------------------------------------------------------
   for i = #elements, 1, -1 do
      local e = elements[i]
      match e with
      | `If{ cond }    ->
         body = `If{ cond, {body} }
      | `Until{ cond } ->
         body = +{stat: if -{cond} then -{exit()} else -{body} end }
      | `Forin{ ... } | `Fornum{ ... } ->
         table.insert (e, {body}); body=e
      end
   end

   -------------------------------------------------------------------
   -- Change breaks into gotos that escape all loops at once.
   -------------------------------------------------------------------
   local cfg = { stat = { }, expr = { } }
   function cfg.stat.down(x)
      match x with
      | `Break -> x <- exit()
      | `Forin{ ... } | `Fornum{ ... } | `While{ ... } | `Repeat{ ... } ->
         return 'break'
      | _ -> -- pass
      end
   end
   function cfg.expr.down(x) if x.tag=='Function' then return 'break' end end
   walk.stat(cfg, body)

   if exit_label then body = { body, `Label{ exit_label } } end
   return body
end

----------------------------------------------------------------------
-- Front-end:
----------------------------------------------------------------------

mlp.lexer:add 'unless'
mlp.stat:del  'for'
mlp.stat:del  'while'

loop_element_list = gg.list{ loop_element, terminators='do' }

mlp.stat:add{
   'for', mlp.for_header, loop_element_list, 'do', mlp.block, 'end',
   builder = xloop_builder }

mlp.stat:add{
   'while', mlp.expr, loop_element_list, 'do', mlp.block, 'end',
   builder = |x| xloop_builder{ `While{x[1]}, x[2], x[3] } }

mlp.stat:add{
   'unless', mlp.expr, 'then', mlp.block, 'end',
   builder = |x| +{stat: if not -{x[1]} then -{x[2]} end} }