#lang racket

; a e-tree represents the kinds of program structures that we want to
; manipulate.  Defined recursively:

; atom := a single primative unstructured element, like a number 42 
;         or symbol 'x

; e-tree := atom | ( e-tree * )
; that is, an e-tree is an atom, or a list of 0 or more e-trees

; a path in an e-tree is a list of indexes that navigate to a 
; sub-component of the e-tree.  Navigation via a path is defined
; by the following function, modelled after list-ref

(define (e-tree-ref e-tree path)
  (if (null? path) e-tree
      (e-tree-ref (list-ref e-tree (first path)) (rest path)))
  )

; examples:
(define (test-e-tree-ref)
  (define t1 '(a (b c) (d (e f) (g h (i j)))))
  
  (map displayln (list
    t1
    " "
    (e-tree-ref t1 '())
    (e-tree-ref t1 '(0))
    (e-tree-ref t1 '(2))
    (e-tree-ref t1 '(2 1 1))
    (e-tree-ref t1 '(2 2 2 1))
                  
    )) 
  (void)            
  )

; given an e-tree and a path, we would like to edit the sub e-tree
; at the endpoint of the path.

; first we need a primative for editing an element a list

; given editing function edit-fn, list l, and position n inside l, 
; return a list
; which is l with (list-ref l n) replaced with (fn (list-ref l n))

(define (list-edit edit-fn l n) 
   (let-values 
      ( [ (a b) (split-at l n) ] )  ; note capture a 2 value result
               
      (if (null? b)
          (error (format "edit-list n of ~a >= length l of ~a" n l))
          (append a (list (edit-fn (first b))) (rest b))
          )
      )
  )

; given tree e-tree, and path edit-path, return an e-tree
; which is the original with the subtree t reached by edit-path
; repaced by (edit-fn t)

(define (e-tree-edit edit-fn e-tree edit-path)
  (cond 
    ; empty path edits entire e-tree
    [ (null? edit-path) (edit-fn e-tree) ]
          
    ; for longer path, replace the (first edit-path) position in e-tree with the
    ; result of recursing down the sub e-tree at that position using the rest
    ; of the edit path
    
    [ else (list-edit 
              (lambda (sub-e-tree) (e-tree-edit edit-fn sub-e-tree (rest edit-path)))
              e-tree 
              (first edit-path))
             ]
  ))
  
(define (test-e-tree-edit)
  (define t1 '(a (b c) (d (e f) (g h (i j)))))
  
  (map displayln (list
    t1
    " "
    (e-tree-edit (lambda (s) (rest s)) t1 '())
    (e-tree-edit (lambda (s) 'f1) t1 '(0))
    (e-tree-edit (lambda (s) (list (second s) (first s))) t1 '(1))
    (e-tree-edit (lambda (s) (list 'f2 s) ) t1 '(2 1 1))
    (e-tree-edit (lambda (s) (list 'lambda (list (e-tree-ref s '(1))) s) ) t1 '(2 2 2))
                  
    )) 
  (void)            
  )

; perform the same edit at multiple points in the tree.  You need
; to be careful if the paths intersect in the sense that the endpoint
; of one path is in a sub e-tree below the endpoint of another path.

(define (e-tree-edit-many edit-fn in-tree edit-paths)
  (foldl (lambda (path cur-e-tree) (e-tree-edit edit-fn cur-e-tree path))
         in-tree edit-paths)
  )

(define (test-e-tree-edit-many)
  (define t1 '(a (b c) (d (e f) (g h (i j)))))
  
  (map displayln (list
    t1
    " "
    (e-tree-edit-many (lambda (s) 'f1) t1 '( (0) (1 1)))
    (e-tree-edit-many (lambda (s) (list 'f2 s) ) t1 '( (2 1 1) (2 2 2)) )
                  
    )) 
  (void)            
  )

; find paths to all instances of term in expression expr
; A path is a list of 0-origin indexes into the expression tree
; indicating how to navigate down to the root to locate the term

(define (e-tree-find-paths target-e-tree e-tree [paths '()] [rootpath '()] )
   (cond
        ; are we at a matching term?
        [ (equal? target-e-tree e-tree) (append paths (list rootpath) ) ]
        
        ; ok, no match, no recursion since e-tree empty
        [ (null? e-tree) paths ]
        
        ; recurse into sub e-trees accumulating paths to matches
        [ (list? e-tree) 
            (second 
               (foldl 
                ; (e s) is current sub e-tree e, prior state s
                ; s state is (n p) where 
                ;    n is position of e in containing e-tree, 
                ;    p collection of paths found so far
                (lambda (e s)
                    (let ( 
                          ; pull apart state
                          (n (first s)) 
                          (p (second s)) 
                          )
                        ; new state - next position, add newly found paths
                        (list
                           (add1 n) 
                           (e-tree-find-paths target-e-tree e p (append rootpath (list n)))
                           )
                    ) 
                )
                (list 0 paths) e-tree
                )
            ) ]
        
        [ else paths ]
    )
) 

(define (test-e-tree-find-paths)
  (define t1 '(a (b a) (b (a f) (f h (b a)))))
  
  (map displayln (list
    t1
    " "
    (e-tree-find-paths 'a t1)
    (e-tree-find-paths 'b t1)
    (e-tree-find-paths '(b a) t1)
    
    
    (e-tree-edit-many (lambda (s) 'delta) t1 (e-tree-find-paths 'a t1))
    (e-tree-edit-many (lambda (s) 'delta) t1 (e-tree-find-paths '(b a) t1) )
    )) 
  (void)            
  )

; NOTE: write a version of e-tree-find-paths that has a test predicate 
; other than equal?

; Program transformations

; Here is an attempt to extract out the continuation for a term
; inside a e-tree.

; assuming that e-tree contains non special forms, then extract the 
; continuation context about the term at point path.


(define (extract-continuation-1 e-tree path)
  (cond 
     ; if the path has length 1, we are at a term ( ... )
    
     ; in first position of list, it is a fn call
     [ (equal? path '(0)) '(lambda (x) x) ]
     
     ; in other positions it is a term to evaluate then pass up to 
     ; enclosing expressions
     [ (null? (rest path)) 
       (let ( [ var (gensym 'v- ) ] )
          (list 'lambda (list var) (list-edit (lambda (s) var) e-tree (first path)) )
          )]
     
     ; the path descends, so we need to capture the context around and
     ; descend further
     
     [ else        
       (let ( [ var (gensym 'v- ) ] )
          (list 'lambda (list var) (e-tree-edit (lambda (s) var)  e-tree path))
          )]
 
         
  ))

(define (f-test x y) 
      (if (null? x) y (cons (first x) (f-test (rest x) (cons (first x) y))) )
      )

; run this test by doing
; (eval (test-convert))
; (test-cc '(a b) '(1 2))
; you should get: '(a b b a 1 2)
; 
(define (test-convert) 
  (define e1 '
    (define (f x y) 
      (if (null? x) y (cons (first x) (f (rest x) (cons (first x) y))) )
    ))
    
  (eval (convert-to-cc-style e1 'test-cc))
  
  )

; converts a cannonical recursive definition into continuation-
; passing style.
; the function MUST have the form
; (define (f ... ) 
;  (if (basecase? ...) base-expression recursion-expression) )
;
(define (convert-to-cc-style defn new-fn-name)
  
  (or (eq? 'if (e-tree-ref defn '(2 0)))
          (error (format "not a if-style definition ~a" defn)))
  
  ; get the parts of the definition
  (let (
        [ args (e-tree-ref defn '(1)) ]
        [ fn-name (e-tree-ref defn '(1 0)) ]
        [ base-test (e-tree-ref defn '(2 1)) ]
        [ base-case (e-tree-ref defn '(2 2)) ]
        [ recursion (e-tree-ref defn '(2 3)) ]
        )
    
    
  ; locate the recursive call in the recursion
  (define p1 (e-tree-find-paths fn-name recursion))
  
  ; make sure only one call
  (or (= 1 (length p1)) 
      (error (format "not a single recursion in ~a" defn) ) )
  
  ; move up to expression that needs result of recursive call
  ; and make that the continuation
  (define cc1 (list-edit 
               (lambda (s) (list 'cc s))
               (extract-continuation-1 recursion (drop-right (first p1) 1))
               2))
  (displayln cc1)
  
  ; insert continuation into the original recursion
  (define new-rec (append (e-tree-ref recursion (drop-right (first p1) 1)) (list cc1)))
 
  (displayln new-rec)
  
  ; create the default continuation
  
  ; rewrite defintion to be in continuation passing style
  (define new-def
    (list 'define (append (second defn) '( (cc (lambda (x) x)))) 
          (list 'if base-test (list 'cc base-case) new-rec) )    
    )
  
 
  (displayln new-def)
  
  ; rename the original function to the new name
  (define final-def 
     (e-tree-edit-many (lambda (s) new-fn-name) new-def (e-tree-find-paths fn-name new-def)) )
    
  (displayln final-def)
 
  ; then eval the final result to create the new function.
  
  final-def
  
  ))