#lang racket/base

(provide
  tree tree? tree-key tree-value tree-children
  make-tree tree-ref tree-set!
  tree-node tree-child tree->list
  tree-for-each tree-find-path)

; Tree structures, with key and possible value on nodes
; Like a table, with a list of keys as reference to a value
(struct tree
  (key ; Key associated to a node. #f for the root of the tree
   (value #:mutable) ; Value associated to a node. May be #f if the node has no value associated with (it should have children then)
   (children #:mutable) ; Children list of a node
  ))

; Constructor
; New node of a tree.
(define (make-tree [key #f])
  (tree key #f (list)))

; Get a child from its key
(define (tree-child t key)
  (let ((child-list
          (filter
            (lambda (x) (equal? (tree-key x) key))
            (tree-children t))))
    (if (null? child-list)
        #f
        (car child-list))))

; Set a value on the given path of keys in the tree
(define (tree-set! tr path value)
  (if (null? path)
    (set-tree-value! tr value)
    (let ((child (tree-child tr (car path))))
      (if child
          (tree-set! child (cdr path) value)
          (let ((new-child (make-tree (car path))))
            (set-tree-children! tr (cons new-child (tree-children tr)))
            (tree-set! new-child (cdr path) value))))))

; Get a sub-tree from its path, or #f if not found
(define (tree-node tr path)
  (if (null? path)
    tr
    (let ((child (tree-child tr (car path))))
      (if child
          (tree-node child (cdr path))
          #f))))

; Get a value from its path, or #f if not found
(define (tree-ref tr path)
  (let ((node (tree-node tr path)))
    (and node (tree-value node))))

; Get the list of all values along a path, or #f if the path does not exist
(define (tree->list tr path [ret (list)])
  (if (null? path)
      (reverse (cons (tree-value tr) ret))
      (let ((child (tree-child tr (car path))))
        (if child
            (tree->list child (cdr path) (cons (tree-value tr) ret))
            #f))))

; Apply the given function on every node of the tree
(define (tree-for-each tr fun)
  (fun (tree-value tr))
  (for-each
    (lambda (c) (tree-for-each c fun))
    (tree-children tr)))

; Get the path to the first item whose value checks the predicate
(define (tree-find-path tr pred (path (list)))
  (if (pred (tree-value tr))
      (cdr (reverse (cons (tree-key tr) path))) ; cdr as the root node key should not be present
      (findf
        (lambda (x) x)
        (map
          (lambda (c) (tree-find-path c pred (cons (tree-key tr) path)))
          (tree-children tr)))))
