#lang racket
(require "parse.rkt"
         (except-in racket/match ==))

(provide simplify)

;; flatten/label-block : list-of-decl list-of-stmt -> block-stmt
;; Desugars `for', converts `if' so that it's always of the form
;; `if <test> then goto <label> else goto <label>', flattens
;; compound statements into the enclosing block, and gives every
;; statement exactly one label. The result usually has lots of
;; "dummy" statements that could easily be eliminated by merging
;; labels.
(define (flatten/label-block decls statements ->stx)
  (define extra-decls null)
  (define new-statements
    (let loop ([l statements])
      (if (null? l)
          null
          (match (car l)
            [(a60:block decls statements)
             (cons (cons (gensym 'block) (flatten/label-block decls statements ->stx))
                   (loop (cdr l)))]
            [(a60:compound statements)
             (loop (append statements (cdr l)))]
            [(a60:branch test then else)
             (if (and (a60:goto? then) (a60:goto? else))
                 (cons (cons (gensym 'branch) (car l))
                       (loop (cdr l)))
                 (let ([then-label (gensym 'then)]
                       [else-label (gensym 'else)]
                       [cont-label (gensym 'if-cont)])
                   (loop
                    (list*
                     (make-a60:branch test (make-a60:goto then-label) (make-a60:goto else-label))
                     (make-a60:label then-label then)
                     (make-a60:goto cont-label)
                     (make-a60:label else-label else)
                     (make-a60:label cont-label (make-a60:dummy))
                     (cdr l)))))]
            [(a60:for variable val-exprs body)
             (let ([body-label (gensym 'for-body)]
                   [cont-label (gensym 'for-cont)])
               (letrec ([make-init+test+increment+loop
                         (lambda (value)
                           (match value
                             [(a60:for-number value)
                              (values (make-a60:assign (list variable) (make-a60:binary 'num 'num
                                                                                        (->stx '+) 
                                                                                        (->stx '0)
                                                                                        value)) ; +0 => number
                                      (->stx #t)
                                      (make-a60:dummy)
                                      #f)]
                             [(a60:for-step start step end)
                              (values (make-a60:assign (list variable) start)
                                      (make-a60:binary 'bool 'num
                                                       (->stx '<=)
                                                       (make-a60:binary 'num 'num
                                                                        (->stx '*)
                                                                        (make-a60:binary 'num 'num (->stx '-) variable end)
                                                                        (make-a60:app (->stx 'sign) (list step)))
                                                       (->stx '0))
                                      (make-a60:assign (list variable) (make-a60:binary 'num 'num (->stx '+) variable step))
                                      #t)]
                             [(a60:for-while value test)
                              (values (make-a60:assign (list variable) value)
                                      test
                                      (make-a60:assign (list variable) value)
                                      #t)]))])
                 (if (= 1 (length val-exprs))
                     (let-values ([(init test inc loop?) (make-init+test+increment+loop (car val-exprs))])
                       (loop (list*
                              init
                              (make-a60:label body-label (make-a60:dummy))
                              (make-a60:branch test 
                                               (make-a60:compound
                                                (list
                                                 body
                                                 inc
                                                 (if loop?
                                                     (make-a60:goto body-label)
                                                     (make-a60:dummy))))
                                               (make-a60:dummy))
                              (cdr l))))
                     (let* ([stage-name (datum->syntax #f (gensym 'stage-number))]
                            [switch-name (datum->syntax #f (gensym 'stage-switch))]
                            [end-switch-name (datum->syntax #f (gensym 'stage-switch))]
                            [stage-var (make-a60:variable stage-name null)]
                            [start-labels (map (lambda (x) (gensym 'stage)) (append val-exprs (list 'extra)))]
                            [end-labels (map (lambda (x) (gensym 'stage)) val-exprs)])
                       (set! extra-decls (list* stage-name 
                                                (cons switch-name start-labels)
                                                (cons end-switch-name end-labels)
                                                extra-decls))
                       (loop
                        (append
                         (list (make-a60:assign (list stage-var) (->stx '0)))
                         (let loop ([start-labels start-labels][end-labels end-labels][val-exprs val-exprs])
                           (if (null? val-exprs)
                               (list (make-a60:label (car start-labels) (make-a60:dummy)))
                               (let-values ([(init test inc loop?) (make-init+test+increment+loop (car val-exprs))])
                                 (list*
                                  (make-a60:label (car start-labels) (make-a60:dummy))
                                  init
                                  (make-a60:branch test 
                                                   (make-a60:goto body-label)
                                                   (make-a60:compound                                                   
                                                    (list
                                                     (make-a60:assign (list stage-var) (make-a60:binary 'num 'num
                                                                                                        (->stx '+)
                                                                                                        (->stx '1)
                                                                                                        stage-var))
                                                     (make-a60:goto (make-a60:subscript switch-name stage-var)))))
                                  (make-a60:label (car end-labels) (make-a60:dummy))
                                  inc
                                  (if loop?
                                      (make-a60:goto (car start-labels))
                                      (make-a60:goto (cadr start-labels)))
                                  (loop (cdr start-labels)
                                        (cdr end-labels)
                                        (cdr val-exprs))))))
                         (list 
                          (make-a60:goto cont-label)
                          (make-a60:label body-label (make-a60:dummy))
                          body
                          (make-a60:goto (make-a60:subscript end-switch-name stage-var))
                          (make-a60:label cont-label (make-a60:dummy)))
                         (cdr l)))))))]
            [(a60:label name statement)
             (cons (cons name (make-a60:dummy))
                   (loop (cons statement (cdr l))))]
            [else
             (cons (cons (gensym 'other) (car l))
                   (loop (cdr l)))]))))
  (make-a60:block
   (append
    (map (lambda (decl)
           (match decl
             [(a60:proc-decl result-type var arg-vars by-value-vars arg-specs body)
              (make-a60:proc-decl result-type var arg-vars by-value-vars arg-specs
                                  (simplify-statement body ->stx))]
             [else decl]))
         decls)
    (map (lambda (extra)
           (if (identifier? extra)
               (make-a60:type-decl (->stx 'integer) (list extra))
               (make-a60:switch-decl (car extra) (map (lambda (x) 
                                                        (make-a60:variable (datum->syntax #f x) null)) 
                                                      (cdr extra)))))
         extra-decls))
   (if (null? new-statements)
       (list (cons (gensym 'other) (make-a60:dummy)))
       new-statements)))

(define (simplify stmt ctx)
  (simplify-statement stmt (lambda (x)
                             (datum->syntax
                              ctx
                              x))))

(define (simplify-statement stmt ->stx)
  (match stmt
    [(a60:block decls statements)
     (flatten/label-block decls statements ->stx)]
    [(a60:compound statements)
     (flatten/label-block null statements ->stx)]
    [else stmt]))
