#!/usr/bin/luajit
--[[ Really simple (and bogus) JIT compiler for a tiny Lisp.

This Lisp has just enough power to express dumb Fibonacci:

    (letrec (fib (lambda (n)
                   (if (< n 2)
                       1
                       (+ (fib (- n 1)) (fib (- n 2))))))
      (fib 40))

It works by compiling the Lisp S-expression to Lua and using the
LuaJIT compiler to compile it to machine code and execute it.  In this
case, starting up LuaJIT, parsing, and generating the Lua takes about
10 milliseconds; executing it, producing the value 165580141, takes
about 2600 milliseconds.  This is roughly 2.8 times the execution time
of a C program <http://canonical.org/~kragen/sw/dev3/fib.c> that does
the same stupid thing.

This compiler is bogus in a few different ways:

1. I haven’t implemented the basic Lisp operations of CONS, CDR, CAR,
   ATOM, NULL, and QUOTE, so there’s no way to do any list processing in
   this “Lisp”.  It’s more of a Mocklisp than a Lisp.

2. Correctly generating Lua code requires attention to when you’re
   compiling a statement and when you’re compiling an expression,
   although function calls are valid in both contexts.  This compiler
   does not pay such attention and consequently generates
   syntactically invalid Lua in many cases.  For example, `(letrec (f
   (lambda () (if 1 (+ 3 4) 1))) (f 5 6 8))` does work, but simpler
   variants such as `(+ 3 4)` do not.  Safe code generation would also
   require some attention to Lua identifier syntax.

3. Not only does this compiler inherit its eagerness, argument
   evaluation order, garbage collection, error handling (including
   dynamic type checking), and JIT compilation from the host Lua
   implementation, it actually inherits scoping, including closures,
   from the host implementation too.

   (Dynamic type checking?  In MY compiler?  It’s more likely than you think:
   `(print (+ (lambda () (print 4)) 4))`)

4. The parse error message “(string expected, got nil)” leaves a lot
   to be desired.  I think I’ve at least removed the stack-overflow
   cases.

Nevertheless, it’s incredibly fucking cool that with LuaJIT you can
write a compiler for a new programming language in two hours and get
performance that beats unoptimized C, even if it would take a day or
more to get it to work adequately.

--]]

---- Reader

-- These read functions return pairs: parsed_value, next_char.
-- Either or both may be nil, meaning respectively an empty list
-- and no next char.  They all take a non-optional first char and a
-- function to get more chars.

function read_sexp(c, getc)
   while c:match("[%s]") do c = getc() end
   if c == '(' then return read_list(getc(), getc) end
   return read_atom(c, getc)
end

function read_list(c, getc)
   while c ~= nil and c:match("[%s]") do c = getc() end -- XXX should not check nil
   if c == ')' then return nil end

   local car, c2 = read_sexp(c, getc)
   if c2 == nil then c2 = getc() end

   return {car=car, cdr=read_list(c2, getc)}
end

function read_atom(c, getc)
   local name = {}

   while c ~= nil and c:match("[^%s()]") do
      table.insert(name, c)
      c = getc()
   end

   name = table.concat(name)
   -- XXX negative numbers
   if not name:match("[^%d.]") then return tonumber(name), c end
   return name, c
end

---- Printer

-- The name of this function is somewhat optimistic.

function pprint(sexp, output)
   if sexp == nil then return output("()") -- XXX unnecessary case
   elseif type(sexp) == 'string' then return output(":" .. sexp) -- XXX round trip fail
   elseif type(sexp) == 'number' then return output('' .. sexp)
   else
      output("(")
      while sexp ~= nil do
         pprint(sexp.car, output)
         sexp = sexp.cdr
         if sexp ~= nil then output(" ") end
      end
      return output(")")
   end
end   

---- Compiler

function consed_values(sexp)  -- Helper for converting to Lua lists
   if sexp == nil then return end
   return sexp.car, consed_values(sexp.cdr)
end

function lcompile(sexp)  -- Lua-compile a sexp
   if sexp == nil then return 'nil' end
   if type(sexp) == 'number' then return tostring(sexp) end
   if type(sexp) == 'string' then return sexp end
   local handler = builtin_table[sexp.car]
   if handler ~= nil then return handler(consed_values(sexp.cdr)) end
   return compile_call(sexp.car, consed_values(sexp.cdr))
end

function compile_call(func, ...)  -- Compile a function call
   local arg_exprs = {}
   for _, arg in ipairs({...}) do
      table.insert(arg_exprs, lcompile(arg))
   end
   return lcompile(func) .. '(' .. table.concat(arg_exprs, ', ') .. ')'
end

function binop(op)  -- Create a binary-operation handler
   return function(a, b)
      return '(' .. lcompile(a) .. ') ' .. op .. ' (' .. lcompile(b) .. ')'
   end
end

builtin_table = {
   letrec = function(vars, ...)
      local names, stmts, body = {}, {}, ...  -- XXX adjusting body
      while vars ~= nil do
         table.insert(names, vars.car)
         table.insert(stmts, vars.car .. ' = ' .. lcompile(vars.cdr.car))
         vars = vars.cdr.cdr
      end
      table.insert(stmts, 'return ' .. lcompile(body))

      return 'local ' .. table.concat(names, ', ') .. '\n'
         .. table.concat(stmts, '\n')
   end,

   lambda = function(args, ...)
      local body_stmts = {}
      for _, stmt in ipairs({...}) do
         table.insert(body_stmts, lcompile(stmt))
      end

      return 'function (' .. table.concat({consed_values(args)}, ', ') .. ')\n'
         .. table.concat(body_stmts, '\n') .. '\n'
         .. 'end\n\n'
   end,

   -- XXX this is fairly bogus
   ['if'] = function(condition, consequent, alternate)
      return 'if ' .. lcompile(condition)
        .. ' then return ' .. lcompile(consequent)
        .. ' else return ' .. lcompile(alternate)
        .. ' end'
   end,

   ['<'] = binop('<'),
   ['*'] = binop('*'),
   ['-'] = binop('-'),
   ['+'] = binop('+'),
}

---- Top level

local getc = function() return io.read(1) end
local sexp = read_sexp(getc(), getc)
pprint(sexp, io.write)
print("")
local compiled = lcompile(sexp)
print(compiled)
local loaded = loadstring(compiled)
print(loaded)
print(loaded())

-- time ./terp.lua <<<'(letrec (f (lambda (n m o) (if (< n 2) 1 (* n (f (- n 1)))))) (f 10))'
-- time ./terp.lua <<<'(letrec (fib (lambda (n) (if (< n 2) 1 (+ (fib (- n 1)) (fib (- n 2)))))) (fib 40))'