Improve TCO handling on Common Lisp easy cases

tco.lisp

(defpackage :ssaavedra.tco
  (:documentation
   "Implements some silly optimizations for improved Tail Call
Optimization (TCO) so that when defining a function or a lambda ending
in a `case' or an `if' the enclosing frame can be removed.")
  (:use :common-lisp)
  (:export :tco-body :case-tco :defun-tco))
(in-package :ssaavedra.tco)

(defun tco-test ()
  (print (tco-body '(1 2
		     (let ((X1 (len-array-heap H V)))
			(let ((B (< I X1)))
			  (cl:break)
			  (progn 2
				 (case B ((t) 1) ((nil) 2)))
					;(format t "test")
			  )
			(case B ((t) 1) ((nil) 2))
			(if t 2 3)
			))
		   'FUNNAME)))


(eval-when (:compile-toplevel :execute :load-toplevel)
  (defun tco-handle-form (form fname)
    (let ((operation (car form)))
      (case operation
	((let) (tco-handle-let form fname))
	((let*) (tco-handle-let form fname))
	((progn) (tco-handle-progn form fname))
	((case) (tco-handle-case form fname))
	((if) (tco-handle-if form fname))
	(t (format t "Cannot optimize tail-call to ~a in defun for ~a~%" form fname)
	   (list form)))))

  (defun tco-handle-let (form fname)
    (let ((operation (car form))
	  (assignments (cadr form))
	  (let-body (cddr form)))
      (list (append (list operation assignments) (tco-body let-body fname)))))

  (defun tco-handle-progn (form fname)
    (list (cons 'progn (tco-body (cdr form) fname))))

  (defun tco-handle-case (form fname)
    (let ((condition (cadr form))
	  (cases (cddr form)))
      (list (macroexpand-1 (list 'case-tco fname condition cases)))))

  (defun tco-handle-if (form fname)
    (let ((condition (cadr form))
	  (then (caddr form))
	  (else (cadddr form)))
      (list (macroexpand-1 (list 'if-tco fname condition then else)))))
  
  (defun tco-body (body fname)
    "Returns a body which is better tail-call optimized if the last
instruction is a case form (optionally inside a let or let*, or a progn)."
    (if (cdr body)
	   (cons (car body) (tco-body (cdr body) fname))
	   (let* ((form (car body)))
	     (if (not (consp form)) ;; Maybe this is a last literal/symbol form
		 body
		 ;; Else, this is a cons TCO-able form
		 (tco-handle-form form fname))))))

(defmacro if-tco (fname cond then else)
  "Optimizes an `if' call as a last call of a function by replacing
  the branches with a return-from call so that the compiler can
  absolutely replace the CALL by a JMP to the called function."
  (list 'if cond
	(list 'return-from fname then)
	(list 'return-from fname else)))

(defmacro case-tco (fname cond cases)
  "Optimizes a case call as last call of a function by replacing the
body forms with a `return-from' call so that the compiler can
absolutely replace the CALL by a JMP to the called function."
  (flet ((case-clause-tco (cases &key name)
	   (mapcar (lambda (l)
		     (let ((guard (car l))
			   (tail (cdr l)))
		       (list guard (cons 'return-from (cons name tail)))))
		   cases)))
    (cons 'cl:case (cons cond (case-clause-tco cases :name fname)))))

(defmacro defun-tco (name args &body body)
  "This macro is like a defun but it modifies the defun body so that
it can optimize tail calls inside a case operation. See `case-tco' for
details on how that is implemented."
  `(defun ,name ,args
     ,(tco-body body name)))

(defmacro lambda-tco (args &body body)
  (let ((blockname (gensym)))
    `(lambda ,args
       (block ,blockname
	 ,(tco-body body blockname)))))

(provide 'ssaavedra.tco)