compiler.ms

; Copyright (c) 2017 Chris Double.  All rights reserved.
; BSD 3-Clause License: http://opensource.org/licenses/BSD-3-Clause
;
; Shen Scheme derived soure code is:
; Copyright (c) 2012-2015 Bruno Deferrari.  All rights reserved.

(define (unbound-symbol? maybe-sym scope)
  (and (symbol? maybe-sym)
       (not (memq maybe-sym scope))))

(define *gensym-counter* 0)

(define (gensym prefix)
  (set! *gensym-counter* (+ 1 *gensym-counter*))
  (string->symbol (string-append prefix (number->string *gensym-counter*))))

(define *yields-boolean*
  `(or and
    kl:=
    null? string? vector? number? pair?
    < > >= <= = eq? equal?
    kl:element? kl:symbol? kl:not kl:variable? kl:boolean?
    kl:empty? kl:shen.pvar? kl:tuple?))

(define (yields-boolean? expr)
  (cond
   ((boolean? expr) #t)
   ((pair? expr)
    (or (memq (car expr) *yields-boolean*)
        (and (eq? 'l2r (car expr))
             (yields-boolean? (car (cdr expr))))))
   (else #f)))

(define (force-boolean expr)
  (if (yields-boolean? expr)
      expr
      `(assert-boolean ,expr)))

(define (compile-expression expr scope)
  (define (unbound? maybe-sym)
    (unbound-symbol? maybe-sym scope))

  (define (ce expr . extra-scope)
    (compile-expression expr (append extra-scope scope)))

  (cond
    ((null? expr) '(quote ()))
    ((unbound? expr) (if (and (symbol? expr) (equal? expr (string->symbol ","))) `(string->symbol ,(symbol->string expr)) `(quote ,expr)))
    ((and (pair? expr) (= (length expr) 4) (eq? (car expr) 'let))
     (define var (list-ref expr 1))
     (define value (list-ref expr 2))
     (define body (list-ref expr 3))
     (emit-let var value body (cons var scope)))
    ((and (pair? expr)
          (> (length expr) 1)
          (eq? (car expr) 'cond))
     (emit-cond (cdr expr) scope))
    ((and (pair? expr)
          (= (length expr) 3)
          (eq? (car expr) 'lambda))
     (define var (list-ref expr 1))
     (define body (list-ref expr 2))
     `(lambda (,var) ,(ce body var)))
    ((and (pair? expr)
          (= (length expr) 3)
          (eq? (car expr) 'and))
     (define expr1 (list-ref expr 1))
     (define expr2 (list-ref expr 2))
     `(and ,(force-boolean (ce expr1))
           ,(force-boolean (ce expr2))))
    ((and (pair? expr)
          (= (length expr) 3)
          (eq? (car expr) 'or))
     (define expr1 (list-ref expr 1))
     (define expr2 (list-ref expr 2))
     `(or ,(force-boolean (ce expr1))
           ,(force-boolean (ce expr2))))
    ((and (list? expr)
          (= (length expr) 4)
          (eq? (car expr) 'if))
     (define test (list-ref expr 1))
     (define then (list-ref expr 2))
     (define else (list-ref expr 3))
     `(if ,(force-boolean (ce test))
           ,(ce then)
           ,(ce else)))
    ((and (list? expr)
          (eq? (car expr) 'trap-error)
          (= (length expr) 3)
          (list? (list-ref expr 2))
          (eq? (car (list-ref expr 2)) 'lambda))
     (define expression (list-ref expr 1))
     (define fun (list-ref expr 2))
     (define args (list-ref fun 1))
     (define body (list-ref fun 2))
     `(guard (lambda (,args) ,(ce body args))
        ,(ce expression)))
    ((and (list? expr)
          (eq? (car expr) 'trap-error)
          (= (length expr) 3))
     (define expression (list-ref expr 1))
     (define handler (list-ref expr 2))
     `(guard (lambda (e) (,(ce handler) e))
        ,(ce expression)))
    ((and (list? expr)
          (eq? (car expr) 'do)
          (= (length expr) 3))
     (define expr1 (list-ref expr 1))
     (define expr2 (list-ref expr 2))
     `(begin ,(ce expr1) ,(ce expr2)))
    ((and (list? expr)
          (eq? (car expr) 'freeze)
          (= (length expr) 2))
      (define expr1 (list-ref expr 1))
      `(lambda () ,(ce expr1)))
    ((equal? expr '(fail)) '(quote shen.fail!))
    ((and (pair? expr)
          (eq? (car expr) 'type)
          (= (length expr) 3))
      (define x (list-ref expr 1))
      (define type (list-ref expr 2))
      (ce x))
    ((and (list? expr)
          (eq? (car expr) 'wasp.)
          (= (length expr) 2))
     (car (string->exprs (cadr expr))))
    ((and (list? expr)
          (eq? (car expr) '=)
          (= (length expr) 3))
     (define v1 (list-ref expr 1))
     (define v2 (list-ref expr 2))
     (emit-equality-check v1 v2 scope))
    ((and (list? expr)
          (eq? (car expr) '/)
          (= (length expr) 3))
     (define a (list-ref expr 1))
     (define b (list-ref expr 2))
     `(/ ,(ce a) ,(ce b)))
    ((and (list? expr)
          (eq? (car expr) 'pos)
          (= (length expr) 3))
     (define str (list-ref expr 1))
     (define n (list-ref expr 2))
     `(make-string 1 (string-ref ,(ce str) ,(ce n))))
    ((and (list? expr)
          (eq? (car expr) 'tlstr)
          (= (length expr) 2))
     (define str (list-ref expr 1))
      `(string-tail ,(ce str) 1))
    ((and (list? expr)
          (eq? (car expr) 'cn)
          (= (length expr) 3))
     (define str1 (list-ref expr 1))
     (define str2 (list-ref expr 2))
     `(string-append ,(ce str1) ,(ce str2)))
    ((and (list? expr)
          (eq? (car expr) 'n->string)
          (= (length expr) 2))
     (define n (list-ref expr 1))
     `(make-string 1 ,(ce n)))
    ((and (list? expr)
          (eq? (car expr) 'string->n)
          (= (length expr) 2))
     (define str (list-ref expr 1))
     `(string-ref ,(ce str) 0))
    ((and (list? expr)
          (eq? (car expr) 'absvector)
          (= (length expr) 2))
     (let ((n (list-ref expr 1)))
       `(make-vector ,(ce n) '(quote shen.fail!))))
    ((and (list? expr)
          (eq? (car expr) '<-address)
          (= (length expr) 3))
     (define v (list-ref expr 1))
     (define n (list-ref expr 2))
     `(vector-ref ,(ce v) ,(ce n)))
    ((and (list? expr)
          (eq? (car expr) 'address->)
          (= (length expr) 4))
     (define v (list-ref expr 1))
     (define n (list-ref expr 2))
     (define x (list-ref expr 3))
     `(let ((_tmp ,(ce v)))
       (vector-set! _tmp ,(ce n) ,(ce x))
       _tmp))
    ((list? expr)
      (emit-application (car expr) (cdr expr) scope))
    (else expr)))

(define (emit-let var value body scope)
  `(let ((,var ,(compile-expression value scope)))
     ,(compile-expression body (cons var scope))))

(define (emit-cond clauses scope)
  `(cond ,@(emit-cond-clauses clauses scope)))

(define (emit-cond-clauses clauses scope)
  (cond
    ((null? clauses) '())
    ((and (pair? clauses)
          (pair? (car clauses))
          (= (length (car clauses)) 2))
      (define test (car (car clauses)))
      (define body (cadr (car clauses)))
      (define rest (cdr clauses))
       (let ((compiled-test (compile-expression test scope))
             (compiled-body (compile-expression body scope))
             (compiled-rest (emit-cond-clauses rest scope)))
         `((,(force-boolean compiled-test) ,compiled-body)
           ,@compiled-rest)))))

(define (emit-equality-check v1 v2 scope)
  (define lhs (compile-expression v1 scope))
  (define rhs (compile-expression v2 scope))
  (cond ((or (unbound-symbol? v1 scope)
             (unbound-symbol? v2 scope)
             (equal? '(fail) v1)
             (equal? '(fail) v2))
         `(eq? ,lhs ,rhs))
        ((and (string? lhs) (string? rhs))
         `(string=? ,lhs ,rhs))
        ((and (vector? lhs) (vector? rhs))
         `(vector=? ,lhs ,rhs))
        ((and (number? lhs) (number? rhs))
         `(= ,lhs ,rhs))
        ((or (string? v1) (string? v2))
         `(equal? ,lhs ,rhs))
        ((null? v1) `(null? ,(compile-expression v2 scope)))
        ((null? v2) `(null? ,(compile-expression v1 scope)))
        (else `(kl:= ,lhs ,rhs))))

(define binary-op-mappings
  '((+ . +)
    (- . -)
    (* . *)
    (> . >)
    (< . <)
    (>= . >=)
    (<= . <=)
    (cons . cons)
    (write-byte . write-u8)))

(define (simple-error msg)
  (error 'shen msg))

(define unary-op-mappings
  '((number? . number?)
    (string? . string?)
    (cons? . pair?)
    (absvector? . vector?)
    (simple-error . simple-error)
    (hd . car)
    (tl . cdr)
    (read-byte . read-u8)))

(define (binary-op-mapping op)
  (let ((res (assq op binary-op-mappings)))
    (and res (cdr res))))

(define (unary-op-mapping op)
  (let ((res (assq op unary-op-mappings)))
    (and res (cdr res))))

(define (prefix-op op)
  (define sop (symbol->string op))
  (define opl (string-length sop))
  (if (and (> opl 5)
           (string=? "wasp." (substring sop 0 5)))
      (string->symbol (substring sop 5 (- opl 5)))
      (string->symbol (string-append "kl:" sop))))

(define (emit-application op params scope)
  (define arity (function-arity op))
  (define partial-call? (not (or (= arity -1) (= arity (length params)))))
  (define args (map (lambda (exp) (compile-expression exp scope)) params))
  (cond ((and (<= arity 0) (null? args))
         (cond ((pair? op) `(,(compile-expression op scope)))
               ((unbound-symbol? op scope) `(,(prefix-op op)))
               (else `(,op))))
        (partial-call?
         (nest-call (nest-lambda op arity '()) args))
        ((or (pair? op) (not (unbound-symbol? op scope)))
         (nest-call (compile-expression op scope) args))
        (else
         (cond ((and (= arity 2) (binary-op-mapping op))
                (cons (binary-op-mapping op) args))
               ((and (= arity 1) (unary-op-mapping op))
                (cons (unary-op-mapping op) args))
               (else
                (let ((op (prefix-op op)))
                  (cons op args)))))))

(define (nest-call op args)
  (if (null? args)
      op
      (nest-call (list op (car args)) (cdr args))))

(define (nest-lambda callable arity scope)
  (define (merge-args f arg)
    (if (pair? f)
        (append f (list arg))
        (list f arg)))

  (if (<= arity 0)
      (compile-expression callable scope)
      (let ((aname (gensym "Y")))
        `(lambda (,aname)
           ,(nest-lambda (merge-args callable aname)
                         (- arity 1)
                         (cons aname scope))))))

(define (kl->wasp expr)
  (cond
    ((and (pair? expr) (eq? (car expr) 'defun))
     (define name (list-ref expr 1))
     (define args (list-ref expr 2))
     (define body (list-ref expr 3))
     `(begin
        (register-function-arity (quote ,name) ,(length args))
        (define (,(prefix-op name) ,@args) ,(compile-expression body args))
        (export ,name)
        (quote ,name)))
    (else (compile-expression expr '()))))