(module class-hierarchy mzscheme

  (require (lib "contract.ss")
           (lib "class.ss")
           (lib "etc.ss")
           (prefix srfi1: (lib "1.ss" "srfi")))

  (define-struct class-forest (trees))
  (define-struct class-tree (parent children))
  ;; A ClassForest is (make-class-forest (Listof ClassTree))
  ;; A ClassTree is (make-class-tree Class ClassForest)

  (provide/contract
   
   [class-forest? (any/c . -> . boolean?)]
   [rename new-class-forest make-class-forest (-> class-forest?)]
   [class-forest-insert ([class? class-forest?] [any/c] . opt-> . class-forest?)]
   [class-forest-exists? (class? class-forest? . -> . boolean?)]
   [class-forest-lookup (class? class-forest? any/c . -> . any/c)]
   [class-forest-trees (class-forest? . -> . (listof class-tree?))]
   [class-forest-empty? (class-forest? . -> . boolean?)]
   
   [class-tree? (any/c . -> . boolean?)]
   [rename new-class-tree make-class-tree
           ([class?] [any/c] . opt-> . class-tree?)]
   [class-tree-insert ([class? class-tree?] [any/c] . opt-> . class-tree?)]
   [class-tree-exists? (class? class-tree? . -> . boolean?)]
   [class-tree-lookup (class? class-tree? . -> . any/c)]
   [class-tree-parent (class-tree? . -> . class?)]
   [class-tree-children (class-tree? . -> . class-forest?)]
   [class-tree-singleton? (class-tree? . -> . boolean?)])

  ;; class-forest-empty? : ClassForest -> Boolean
  ;; Reports whether a forest has no trees.
  (define (class-forest-empty? forest)
    (null? (class-forest-trees forest)))

  ;; class-tree-singleton? : ClassTree -> Boolean
  ;; Reports whether a tree has no children.
  (define (class-tree-singleton? tree)
    (class-forest-empty? (class-tree-children tree)))

  ;; class-forest-exists? : Class ClassForest -> Boolean
  ;; Reports whether a forest contains the given class.
  (define (class-forest-exists? class% forest)
    (ormap (lambda (tree) (class-tree-exists? class% tree))
           (class-forest-trees forest)))

  ;; class-tree-exists? : Class ClassTree -> Boolean
  ;; Reports whether a tree contains the given class.
  (define (class-tree-exists? class% tree)
    (let* ([parent% (class-tree-parent tree)]
           [children (class-tree-children tree)])
      (and (subclass? class% parent%)
           (or (subclass? parent% class%)
               (class-forest-exists? class% children)))))

  ;; class-forest-lookup : Class ClassForest [Any] -> Any
  ;; Looks up the value associated with a class in a forest.
  (define class-forest-lookup
    (opt-lambda (class% forest [default #f])
      default))

  ;; class-tree-lookup : Class ClassTree -> Any
  ;; Looks up the value associated with a class in a tree.
  (define (class-tree-lookup class% tree)
    #f)

  ;; new-class-forest : -> ClassForest
  ;; Creates an empty forest of classes.
  (define (new-class-forest)
    (make-class-forest null))

  ;; new-class-tree : Class -> ClassTree
  ;; Creates a class tree with only a root.
  (define new-class-tree
    (opt-lambda (parent [value #f])
      (make-class-tree parent (new-class-forest))))

  ;; class-tree-insert : Class ClassTree -> ClassTree
  ;; Inserts a single class with an optional associated value
  ;; to its proper place in a tree hierarchy.
  (define class-tree-insert
    (opt-lambda (class% tree [value #f])
      (let* ([parent% (class-tree-parent tree)])
        (unless (subclass? class% parent%)
          (error 'class-tree-insert "~s is not a subclass of ~s"
                 class% parent%))
        ;; Using bidirectional subclassing rather than eq? for class equality.
        ;; This works even if there are "wrapped" classes of any sort.
        (if (subclass? parent% class%)
            tree
            (make-class-tree
             parent%
             (class-forest-insert class% (class-tree-children tree)))))))

  ;; class-forest-insert : Class ClassForest [Any] -> ClassForest
  ;; Inserts a single class with an optional associated value
  ;; to its proper place in a forest hierarchy.
  (define class-forest-insert
    (opt-lambda (class% forest [value #f])
      (let* ([trees (class-forest-trees forest)])
        (make-class-forest
         (or (insert-into-existing-tree class% trees)
             (insert-new-tree class% trees))))))

  ;; insert-into-existing-tree
  ;; : Class (Listof ClassTree) -> (Or (Listof ClassTree) #f)
  ;; Inserts a single class to the proper tree, or returns #f if no such tree.
  (define (insert-into-existing-tree class% trees)
    (recur search ([done null]
                   [todo trees])
      (if (null? todo)
          #f
          (let* ([tree (car todo)]
                 [rest (cdr todo)])
            (if (subclass? class% (class-tree-parent tree))
                (srfi1:append-reverse
                 done (cons (class-tree-insert class% tree) rest))
                (search (cons tree done) rest))))))

  ;; insert-new-tree : Class (Listof ClassTree) -> (Listof ClassTree)
  ;; Creates a new tree for a class, and adds any existing trees to its
  ;; hierarchy if necessary.
  (define (insert-new-tree class% trees)
    (let*-values ([(children peers)
                   (srfi1:partition
                    (lambda (tree)
                      (subclass? (class-tree-parent tree) class%))
                    trees)])
      (append peers
              (list (make-class-tree class% (make-class-forest children))))))

  )