advent-of-code-2022/day-11/peg-tree-utils.scm

(library (peg-tree-utils)
  (export find-in-tree
          find-in-tree*
          tree-refs)
  (import
   (except (rnrs base) let-values)
   (only (guile) lambda* λ
         current-output-port
         simple-format)
   (srfi srfi-1))


  (define find-in-tree
    (λ (peg-tree label equal-test)
      (define traverse
        (λ (peg-tree label)
          (cond
           [(null? peg-tree) #f]
           [(pair? (car peg-tree))
            (or (traverse (first peg-tree) label)
                (traverse (cdr peg-tree) label))]
           [(symbol? (car peg-tree))
            (cond
             [(equal-test (car peg-tree) label)
              (car (cdr peg-tree))]
             [else
              (traverse (cdr peg-tree) label)])]
           [else
            (traverse (cdr peg-tree) label)])))
      (traverse peg-tree label)))


  (define tree-refs
    (lambda* (peg-tree refs #:key (equal-test eq?))
      (let traverse ([tree° peg-tree]
                     [refs° refs])
        (cond
         [(null? refs°) tree°]
         [(null? tree°) #f]
         [(pair? (car tree°))
          (or (traverse (car tree°) refs°)
              (traverse (cdr tree°) refs°))]
         [(symbol? (car tree°))
          (cond
           ;; success case
           [(equal-test (car tree°) (car refs°))
            (traverse (cdr tree°) (cdr refs°))]
           [else
            (traverse (cdr tree°) refs°)])]
         [else
          (traverse (cdr tree°) refs°)]))))


  (define find-in-tree*
    (λ (peg-tree filter-proc)

      (define traverse
        (λ (subtree cont)
          ;; (simple-format (current-output-port)
          ;;                "working with subtree ~a\n"
          ;;                subtree)
          (cond
           [(null? subtree) (cont)]
           [(pair? (first subtree))
            (traverse (first subtree)
                      (λ () (traverse (cdr subtree) cont)))]
           [(filter-proc (first subtree))
            (cons subtree (cont))]
           [else
            (traverse (cdr subtree) cont)])))

      (traverse peg-tree (λ () '())))))