;; Iterating over trees

(define (tree-func? t s)
  (or (== s (tree-get-label t))
      (and (in? (tree-get-label t) '("expand" "var_expand" "hide_expand"))
           (== s (tree->string (tree-ref t 0))))))

(define (tree-compound-arity t)
  (if (tree-atomic? t) 0 (tree-arity t)))

(define (search-in-tree t label)
  (let down ((t t) (ip '()))
    (if (tree-func? t label)
        (reverse ip)
        (let right ((i 0))
          (and (< i (tree-compound-arity t))
               (or (down (tree-ref t i) (cons i ip))
		   (right (1+ i))))))))

(define (safe-tree-ref t i)
  (if (< i (tree-compound-arity t))
      (tree-ref t i)
      (error "safe-tree-ref, index out of range")))

(define (subtrees-on-path t p)
  ;; Stack of all the subtrees traversed when getting (subtree t p).
  ;; (subtree t p) is the first item, and t is the last item.
  ;; This is useful to initialize the backtracking stack for tree searches.
  (define (kons i ts)
    (cons (safe-tree-ref (first ts) i) ts))
  (list-fold kons (list t) p))

(define (search-in-tree-from/down label proc ts ip t)
  (if (tree-func? t label)
      (or (proc (reverse ip) t)
	  (search-in-tree-from/up label proc ts ip))
      (search-in-tree-from/right label proc ts ip t 0)))

(define (search-in-tree-from/up label proc t+ts i+ip)
  (and (pair? t+ts)
       (search-in-tree-from/right
	label proc (cdr t+ts) (cdr i+ip) (car t+ts) (1+ (car i+ip)))))

(define (search-in-tree-from/right label proc ts ip t i)
  (if (< i (tree-compound-arity t))
      (search-in-tree-from/down
       label proc (cons t ts) (cons i ip) (tree-ref t i))
      (search-in-tree-from/up label proc ts ip)))

(define (search-in-tree-from t path label proc)
  (if (null? path)
      (search-in-tree-from/down label proc '() '() t)
      (let ((t+ts (subtrees-on-path t path)))
	(search-in-tree-from/down
	 label proc (cdr t+ts) (reverse path) (car t+ts)))))

;(set-trace-level! search-in-tree-from search-in-tree-from/down
;		  search-in-tree-from/up search-in-tree-from/right)

; (define (display-found-em p+t)
;   (display* "Found em at: " (first p+t) "\n")
;   (display* "Found em data: " (tree->object (second p+t)) "\n"))

; (define (find-next-em)
;   (let ((p+t (search-in-tree-from
; 	      (the-buffer) (but-last (the-path)) "em" list)))
;     (tm-go-to (rcons (first p+t) 0))
;     (display-found-em p+t)))

; (define (find-all-em)
;   (define found '())
;   (define (proc p t) (set-rcons! found (list p t)) #f)
;   (search-in-tree-from
;    (the-buffer) (but-last (the-path)) "em" proc)
;   (for-each display-found-em found))

; (kbd-map
;   ("C-x e 2" (backtrack (find-next-em)))
;   ("C-x e 1" (backtrack (find-all-em))))