#lang racket

;; A tree-of-X is
;;  - #f
;;  - (node tree-or-X X tree-of-X)
(struct node (l v r))

(define t1 (node #f
                 "apple"
                 (node (node #f "banana"  #f)
                       "coconut"
                       (node (node #f "donut" #f)
                             "eggplant"
                             (node #f "fig" #f)))))

(define t2 (node (node #f "banana"  #f)
                 42
                 (node #f "donut" #f)))

;; ----------------------------------------
;; Strightforward count of regexp matchs in the tree:

;; count : tree-of-string regexp -> number
(define (count t rx)
  (cond
    [(not t) 0]
    [else (+ (if (regexp-match? rx (node-v t)) 1 0)
             (count (node-l t) rx)
             (count (node-r t) rx))]))

'count
(count t1 #rx"a")

;; We might have many more functions like `count`...

;; ----------------------------------------
;; Suppose that extracting the string from a node costs
;; something (e.g., remote database lookup):

;; extract-string : tree-of-string -> (values string number)
(define (extract-string t)
  (values (node-v t) 1))

;; count/cost : tree-of-string regexp -> (values number number)
(define (count/cost t rx)
  (cond
    [(not t) (values 0 0)]
    [else 
     (define-values (cnt1 cst1) (count/cost (node-l t) rx))
     (define-values (cnt2 cst2) (count/cost (node-r t) rx))
     (define-values (s cst) (extract-string t))
     (values (+ (if (regexp-match? rx s) 1 0)
                cnt1
                cnt2)
             (+ cst cst1 cst2))]))

'count/cost
(count/cost t1 #rx"a")

;; All other functions functions like `count` must use the same
;; pattern...

;; ----------------------------------------
;; Suppose that we want to fail at the point that we discover
;; a non-string in the tree --- simulating communication failure,
;; for example.

;; extract-string/fail : tree-of-X -> string-or-#f
(define (extract-string/fail t)
  (if (string? (node-v t))
      (node-v t)
      #f))

;; count/fail : tree-of-X regexp -> number-or-#f
(define (count/fail t rx)
  (cond
    [(not t) 0]
    [else 
     (define cnt1 (count/fail (node-l t) rx))
     (cond
       [(not cnt1) #f]
       [else
        (define s (extract-string/fail t))
        (cond
          [(not s) #f]
          [else
           (define cnt2 (count/fail (node-r t) rx))
           (cond
             [(not cnt2) #f]
             [else
              (+ (if (regexp-match? rx (node-v t)) 1 0)
                 cnt1
                 cnt2)])])])]))

'count/fail
(count/fail t1 #rx"a")
(count/fail t2 #rx"a")

;; ----------------------------------------

;; Combine both costs and failure!

;; extract-string/c+f : tree-of-X -> (values string-or-#f number)
(define (extract-string/c+f t)
  (if (string? (node-v t))
      (values (node-v t) 1)
      (values #f 0)))

;; count/c+f : tree-of-X regexp -> (values number-or-#f number)
(define (count/c+f t rx)
  (cond
    [(not t) (values 0 0)]
    [else 
     (define-values (cnt1 cst1) (count/c+f (node-l t) rx))
     (cond
       [(not cnt1) (values #f cst1)]
       [else
        (define-values (s cst) (extract-string/c+f t))
        (cond
          [(not s) (values #f (+ cst1 cst))]
          [else
           (define-values (cnt2 cst2) (count/c+f (node-r t) rx))
           (cond
             [(not cnt2) (values #f (+ cst cst1 cst2))]
             [else
              (values (+ (if (regexp-match? rx (node-v t)) 1 0)
                         cnt1
                         cnt2)
                      (+ cst cst1 cst2))])])])]))

'count/c+f
(count/c+f t1 #rx"a")
(count/c+f t2 #rx"a")

;; ----------------------------------------
;; First try at abstracting over the pattern ---
;; handles failure, but doesn't accumulate cost:

(module+ failed-attempt
  
  (struct priced (v cost) #:transparent)
  
  ;; return : number-or-string -> priced
  (define (return v)
    (priced v 0))
  
  ;; extract : tree-of-X -> priced
  (define (extract t)
    (if (string? (node-v t))
        (priced (node-v t) 1)
        (priced #f 0)))
  
  ;; bind : priced (number -> priced) -> priced
  (define (bind p k)
    (cond
      [(not (priced-v p)) p]
      [else (k (priced-v p))]))
  
  ;; count/monad : tree-of-X regexp -> priced
  (define (count/monad t rx)
    (cond
      [(not t) (return 0)]
      [else
       (bind (count/monad (node-l t) rx)
             (lambda (cnt1)
               (bind (extract t)
                     (lambda (s)
                       (define cnt (if (regexp-match? rx (node-v t)) 1 0))
                       (bind (count/monad (node-r t) rx)
                             (lambda (cnt2)
                               (return (+ cnt1 cnt cnt2))))))))]))
  
  'count/monad
  (count/monad t1 #rx"a")
  (count/monad t2 #rx"a"))
  
;; ----------------------------------------
;; Successful monadic implementation:

;; Part of the "internal" representation:
(struct priced (v cost))

;; M = (number -> priced)

;; return : number -> M
(define (return v)
  (lambda (cst)
    (priced v cst)))

;; bind : M (number -> M) -> M
(define (bind proc k)
  (lambda (cst)
    (define p (proc cst))
    (cond
      [(not (priced-v p)) p]
      [else ((k (priced-v p)) (priced-cost p))])))

;; The above two operations are the standard monad operations,
;; but many monad implementations will provide additional
;; operators:

;; extract : tree-of-X -> M
(define (extract t)
  (if (string? (node-v t))
      (lambda (cst)
        (priced (node-v t) (add1 cst)))
      (lambda (cst)
        (priced #f cst))))

;; A monad implementation will also normally have
;; a `run` operation to get things going and extract
;; the result:

;; run : M -> (values number-or-#f number)
(define (run proc)
  (define p (proc 0))
  (values (priced-v p)
          (priced-cost p)))

;; count/monad : tree-of-X regexp -> M
(define (count/monad t rx)
  (cond
    [(not t) (return 0)]
    [else
     (bind (count/monad (node-l t) rx)
           (lambda (cnt1)
             (bind (extract t)
                   (lambda (s)
                     (bind (count/monad (node-r t) rx)
                           (lambda (cnt2)
                             (return (+ (if (regexp-match? rx s) 1 0)
                                        cnt1
                                        cnt2))))))))]))

'count/monad
(run (count/monad t1 #rx"a"))
(run (count/monad t2 #rx"a"))

;; ----------------------------------------
;; `do` syntax makes monad style more readable:

(define-syntax do
  (syntax-rules (<-)
    [(do () body) body]
    [(do ([id0 <- rhs0] [id <- rhs] ...) body)
     (bind rhs0
           (lambda (id0)
             (do ([id <- rhs] ...) body)))]))

;; count/do : tree-of-X regexp -> M
(define (count/do t rx)
  (cond
    [(not t) (return 0)]
    [else
     (do ([cnt1 <- (count/monad (node-l t) rx)]
          [s <- (extract t)]
          [cnt2 <- (count/monad (node-r t) rx)])
       (return (+ (if (regexp-match? rx s) 1 0)
                  cnt1
                  cnt2)))]))

'count/do
(run (count/do t1 #rx"a"))
(run (count/do t2 #rx"a"))