test-phog.scm

(load "gram.scm")
(load "tree.scm")
(load "phog.scm")

;; ============================================================
;; Test harness
;; ============================================================

(define *test-pass* 0)
(define *test-fail* 0)

(define (test name expected actual)
  (if (equal? expected actual)
      (begin
        (set! *test-pass* (+ *test-pass* 1))
        (printf "  ✓ ~a\n" name))
      (begin
        (set! *test-fail* (+ *test-fail* 1))
        (printf "  ✗ ~a\n    expected: ~a\n    actual:   ~a\n" name expected actual))))

(define (test-summary)
  (printf "\n~a passed, ~a failed\n" *test-pass* *test-fail*))

;; ============================================================
;; Test corpus: simple arithmetic expressions
;; ============================================================
;;
;; Grammar:
;;   expr → num op num
;;   expr → num
;;   num  → *token* num "43" etc.
;;   op   → *token* plus "+" | *token* minus "-"

(define corpus
  (list
    ;; "43+2"
    '(expr (num "43") (*token* plus "+") (num "2"))
    ;; "1-99"
    '(expr (num "1") (*token* minus "-") (num "99"))
    ;; "7"
    '(expr (num "7"))
    ;; "100+200"
    '(expr (num "100") (*token* plus "+") (num "200"))))

;; ============================================================
;; tree.scm tests
;; ============================================================

(printf "\n=== tree.scm tests ===\n")

(printf "\n--- CST predicates ---\n")

(test "cst-node? on expr"
      #t
      (cst-node? '(expr (num "43") (*token* plus "+") (num "2"))))

(test "cst-node? on token"
      #f
      (cst-node? '(*token* plus "+")))

(test "token? on token"
      #t
      (token? '(*token* plus "+")))

(test "token? on expr"
      #f
      (token? '(expr (num "43"))))

(test "cst-leaf? on string"
      #t
      (cst-leaf? "hello"))

(test "cst-leaf? on node"
      #f
      (cst-leaf? '(num "43")))

(printf "\n--- CST accessors ---\n")

(test "cst-node-name"
      'expr
      (cst-node-name '(expr (num "43") (*token* plus "+") (num "2"))))

(test "cst-node-children length"
      3
      (length (cst-node-children '(expr (num "43") (*token* plus "+") (num "2")))))

(test "token-name"
      'plus
      (token-name '(*token* plus "+")))

(test "token-text"
      "+"
      (token-text '(*token* plus "+")))

(printf "\n--- cst->string ---\n")

(test "cst->string token"
      "+"
      (cst->string '(*token* plus "+")))

(test "cst->string simple node"
      "43"
      (cst->string '(num "43")))

(test "cst->string expr"
      "43+2"
      (cst->string '(expr (num "43") (*token* plus "+") (num "2"))))

(test "cst->string string leaf"
      "hello"
      (cst->string "hello"))

(printf "\n--- cst-child-types ---\n")

(test "child types of binary expr"
      '(num plus num)
      (cst-child-types '(expr (num "43") (*token* plus "+") (num "2"))))

(test "child types of unary expr"
      '(num)
      (cst-child-types '(expr (num "7"))))

(test "child types with anonymous"
      '(num *anonymous* num)
      (cst-child-types '(expr (num "1") " + " (num "2"))))

(printf "\n--- cst-collect-types ---\n")

(test "collect types from corpus"
      #t
      (let ((types (cst-collect-types corpus)))
        (and (memq 'expr types)
             (memq 'num types)
             #t)))

(printf "\n--- cst-collect-instances ---\n")

(test "collect expr instances"
      4
      (length (cst-collect-instances corpus 'expr)))

(test "collect num instances"
      7  ; 2+2+1+2 = 7 num nodes total
      (length (cst-collect-instances corpus 'num)))

(printf "\n--- cst-collect-child-patterns ---\n")

(test "expr child patterns"
      #t
      (let ((pats (cst-collect-child-patterns corpus 'expr)))
        (and (member '(num plus num) pats)
             (member '(num minus num) pats)
             (member '(num) pats)
             #t)))

(printf "\n--- cst-walk left-sibling-type ---\n")

;; Verify left-sibling-type is passed correctly
(let ((results '()))
  (cst-walk (list '(expr (num "43") (*token* plus "+") (num "2")))
    (lambda (node parent-type position grandparent-type left-sibling-type)
      (set! results
        (cons (list (cst-node-name node) parent-type position grandparent-type left-sibling-type)
              results))))
  (let ((reversed (reverse results)))
    ;; expr is root: parent=#f, pos=0, gp=#f, lsib=#f
    (test "cst-walk root has #f left-sibling"
          '(expr #f 0 #f #f)
          (car reversed))
    ;; num at position 0 of expr: parent=expr, pos=0, gp=#f, lsib=#f (first child)
    (test "cst-walk first child has #f left-sibling"
          '(num expr 0 #f #f)
          (cadr reversed))
    ;; num at position 2 of expr: parent=expr, pos=2, gp=#f, lsib=plus
    (test "cst-walk third child has plus left-sibling"
          '(num expr 2 #f plus)
          (caddr reversed))))

;; ============================================================
;; phog.scm tests — statistics
;; ============================================================

(printf "\n=== phog.scm tests ===\n")

(printf "\n--- alist helpers ---\n")

(test "alist-ref found"
      42
      (alist-ref 'a '((a . 42) (b . 99)) #f))

(test "alist-ref not found"
      #f
      (alist-ref 'c '((a . 42) (b . 99)) #f))

(test "alist-increment new key"
      '((a . 1))
      (alist-increment 'a '()))

(test "alist-increment existing key"
      '((a . 2))
      (alist-increment 'a '((a . 1))))

(printf "\n--- extract-phog-stats ---\n")

(let-values (((phog-table pattern-table) (extract-phog-stats corpus)))
  (test "phog-table is non-empty"
        #t
        (not (null? phog-table)))

  (test "pattern-table is non-empty"
        #t
        (not (null? pattern-table)))

  ;; Check that expr has patterns
  (let ((expr-patterns (alist-ref 'expr pattern-table '())))
    (test "expr has patterns in pattern-table"
          #t
          (not (null? expr-patterns)))

    (test "expr pattern (num plus num) counted"
          #t
          (let ((entry (assoc '(num plus num) expr-patterns)))
            (and entry (> (cdr entry) 0)))))

  ;; Check per-nonterminal PHOG structure
  (test "phog-table has expr entry"
        #t
        (and (assq 'expr phog-table) #t))

  (test "phog-table has num entry"
        #t
        (and (assq 'num phog-table) #t))

  ;; Check PHOG context: expr subtable should have context entries
  (let ((expr-subtable (alist-ref 'expr phog-table '())))
    (test "expr subtable is non-empty"
          #t
          (not (null? expr-subtable)))

    ;; Position 0 of expr with default context (expr 0 #f #f #f) should have num
    (let ((ctx-entry (alist-ref '(expr 0 #f #f #f) expr-subtable '())))
      (test "phog context (expr 0 #f #f #f) has num"
            #t
            (let ((num-entry (assq 'num ctx-entry)))
              (and num-entry (> (cdr num-entry) 0))))))

  (printf "\n--- phog-lookup (Witten-Bell) ---\n")

  ;; phog-lookup now takes 3 args: table, nonterminal, context
  (let ((probs (phog-lookup phog-table 'expr '(expr 0 #f #f #f))))
    (test "phog-lookup returns probabilities"
          #t
          (not (null? probs)))

    (test "num is most frequent at position 0 of expr"
          #t
          (let ((num-prob (alist-ref 'num probs 0)))
            (> num-prob 0))))

  ;; Smoothing: probabilities should sum to ~1
  (let ((probs (phog-lookup phog-table 'expr '(expr 0 #f #f #f))))
    (test "smoothed probabilities sum to ~1"
          #t
          (let ((total (apply + (map cdr probs))))
            (< (abs (- total 1)) 1/1000))))

  ;; Smoothing with unseen context should still return results (backoff)
  (let ((probs (phog-lookup phog-table 'expr '(expr 0 foo bar baz))))
    (test "unseen context returns non-empty via backoff"
          #t
          (not (null? probs)))
    (test "unseen context probabilities sum to ~1"
          #t
          (let ((total (apply + (map cdr probs))))
            (< (abs (- total 1)) 1/1000))))

  ;; Witten-Bell: full context with high counts should have higher lambda
  ;; than sparse context — we verify indirectly: the MLE component should
  ;; have influence when there is data at the full context level.
  (let ((full-probs (phog-lookup phog-table 'expr '(expr 0 #f #f #f)))
        (partial-probs (phog-lookup phog-table 'expr '(expr 0 #f #f))))
    (test "full and partial context both return results"
          #t
          (and (not (null? full-probs))
               (not (null? partial-probs)))))

  ;; Non-existent nonterminal returns empty
  (test "phog-lookup for unknown nonterminal returns empty"
        '()
        (phog-lookup phog-table 'nonexistent '(foo 0 #f #f #f)))

  (printf "\n--- per-nonterminal context functions ---\n")

  ;; Default context function produces 5-element list
  (test "default-context-fn produces 5-element list"
        '(expr 0 stmt #f #f)
        (default-context-fn 'expr 0 'stmt #f #f))

  ;; Custom context function: only use node-type and position
  (register-context-fn! 'test-nt
    (lambda (node-type position parent-type grandparent-type left-sibling-type)
      (list node-type position)))

  (test "custom context-fn returns shorter key"
        '(test-nt 0)
        ((get-context-fn 'test-nt) 'test-nt 0 'parent 'gp 'lsib))

  (test "unregistered nonterminal uses default"
        #t
        (eq? default-context-fn (get-context-fn 'unregistered)))

  ;; Clean up test registration
  (set! *context-fns* '())

  (printf "\n--- condp-static ---\n")

  ;; condp-static should order branches by weight
  (test "condp-static explores highest weight first"
        '(42)
        (run 1 (q)
          (condp-static
            1 (lambda (st) ((== q 'low) st))
            10 (lambda (st) ((== q 42) st))
            5 (lambda (st) ((== q 'mid) st)))))

  (test "condp-static explores all branches"
        '(42 mid low)
        (run* (q)
          (condp-static
            1 (lambda (st) ((== q 'low) st))
            10 (lambda (st) ((== q 42) st))
            5 (lambda (st) ((== q 'mid) st)))))

  (printf "\n--- sym-matcho ---\n")

  ;; Test matching symbol lists against regex-over-symbols
  (test "sym-matcho single sym"
        '(())
        (run* (rest) (sym-matcho '(sym num) '(num) rest '() #f 0 #f)))

  (test "sym-matcho cat"
        '(())
        (run* (rest) (sym-matcho '(cat (sym num) (cat (sym plus) (sym num)))
                                  '(num plus num)
                                  rest '() #f 0 #f)))

  (test "sym-matcho alt"
        '(())
        (run 1 (rest) (sym-matcho '(alt (sym plus) (sym minus))
                                    '(minus)
                                    rest '() #f 0 #f)))

  (test "sym-matcho alt+cat matches both patterns"
        #t
        (let ((regex '(cat (sym num) (cat (alt (sym plus) (sym minus)) (sym num)))))
          (and (not (null? (run 1 (rest) (sym-matcho regex '(num plus num) rest '() #f 0 #f))))
               (not (null? (run 1 (rest) (sym-matcho regex '(num minus num) rest '() #f 0 #f)))))))

  (printf "\n--- synth-rule-for-symbol ---\n")

  ;; Single pattern: should produce a direct concatenation
  (let ((rule (synth-rule-for-symbol 'num '((num)) '() phog-table pattern-table 5)))
    (test "synth single-child rule"
          '(sym num)
          rule))

  (let ((rule (synth-rule-for-symbol 'binary-expr
                '((num plus num))
                '() phog-table pattern-table 5)))
    (test "synth single-pattern rule"
          #t
          ;; Should be a concatenation of sym references
          (not (null? (run 1 (rest)
                        (sym-matcho rule '(num plus num) rest '() #f 0 #f))))))

  ;; Multiple patterns with alternation
  (let ((rule (synth-rule-for-symbol 'binary-op
                '((num plus num) (num minus num))
                '() phog-table pattern-table 10)))
    (test "synth multi-pattern rule matches pattern 1"
          #t
          (not (null? (run 1 (rest)
                        (sym-matcho rule '(num plus num) rest '() #f 0 #f)))))
    (test "synth multi-pattern rule matches pattern 2"
          #t
          (not (null? (run 1 (rest)
                        (sym-matcho rule '(num minus num) rest '() #f 0 #f))))))

  (printf "\n--- learn-grammar ---\n")

  ;; Test the full learning pipeline with the corpus
  (let ((grammar (learn-grammar corpus phog-table 10 pattern-table)))
    (test "grammar is non-empty"
          #t
          (not (null? grammar)))

    (test "grammar has expr rule"
          #t
          (and (assq 'expr grammar) #t))

    (test "grammar has num rule"
          #t
          (and (assq 'num grammar) #t))

    (printf "\nLearned grammar:\n")
    (print-grammar grammar)

    (printf "\n--- verify-grammar ---\n")

    ;; Verify all token-level instances parse correctly
    (let ((results (verify-grammar corpus)))
      (printf "\nVerification results:\n")
      (for-each
        (lambda (r) (printf "  ~a: ~a\n" (car r) (cdr r)))
        results))))

;; ============================================================
;; More complex corpus test: with anonymous strings
;; ============================================================

(printf "\n=== Complex corpus test ===\n")

(define corpus2
  (list
    ;; "let x = 42"
    '(let-stmt (*token* let-kw "let") " " (ident "x") " = " (num "42"))
    ;; "let y = 7"
    '(let-stmt (*token* let-kw "let") " " (ident "y") " = " (num "7"))))

(test "cst->string let-stmt"
      "let x = 42"
      (cst->string (car corpus2)))

(test "child types of let-stmt"
      '(let-kw *anonymous* ident *anonymous* num)
      (cst-child-types (car corpus2)))

(let-values (((phog pattern) (extract-phog-stats corpus2)))
  (test "let-stmt pattern table"
        #t
        (let ((pats (alist-ref 'let-stmt pattern '())))
          (and (not (null? pats))
               ;; Should have exactly one unique pattern
               (= 1 (length pats))))))

;; ============================================================
;; Character-level synthesis tests
;; ============================================================

(printf "\n=== Character-level synthesis tests ===\n")

(printf "\n--- synth-char-level-rule ---\n")

(test "synth-char-level-rule finds rep-digit for numeric texts"
      `(rep ,(digit))
      (synth-char-level-rule '("43" "1" "99" "7" "100")))

(test "synth-char-level-rule finds rep-alpha for letter texts"
      `(rep ,(alpha))
      (synth-char-level-rule '("foo" "Bar" "x" "ABC")))

(test "synth-char-level-rule finds rep-lower for lowercase texts"
      `(rep ,(lower))
      (synth-char-level-rule '("foo" "bar" "x" "abc")))

;; "hello world" matches rep-any since space is in (any) range
(test "synth-char-level-rule matches rep-any for texts with spaces"
      `(rep ,(any))
      (synth-char-level-rule '("hello world" "foo bar")))

(printf "\n--- learn-grammar uses char-level synthesis for num ---\n")

(let-values (((phog-table pattern-table) (extract-phog-stats corpus)))
  (let ((grammar (learn-grammar corpus phog-table 10 pattern-table)))
    (let ((num-rule (cdr (assq 'num grammar))))
      (test "num rule is rep-digit (not alt of literals)"
            `(rep ,(digit))
            num-rule)

      ;; The rep-digit rule should match all original corpus numbers
      (test "rep-digit matches '43'"
            #t
            (parse? 'num "43"))
      (test "rep-digit matches '1'"
            #t
            (parse? 'num "1"))
      (test "rep-digit matches '99'"
            #t
            (parse? 'num "99"))
      ;; And it should generalize to unseen numbers
      (test "rep-digit matches unseen '555'"
            #t
            (parse? 'num "555"))
      (test "rep-digit rejects non-digit 'abc'"
            #f
            (parse? 'num "abc")))))

;; ============================================================
;; Anonymous string handling tests
;; ============================================================

(printf "\n=== Anonymous string handling tests ===\n")

(printf "\n--- collect-anonymous-texts ---\n")

(let ((instances (cst-collect-instances corpus2 'let-stmt)))
  (let ((anon-map (collect-anonymous-texts instances
                    '(let-kw *anonymous* ident *anonymous* num))))
    (test "anonymous position 1 is space"
          " "
          (cdr (assv 1 anon-map)))
    (test "anonymous position 3 is ' = '"
          " = "
          (cdr (assv 3 anon-map)))))

(printf "\n--- learn-grammar with anonymous strings ---\n")

(let-values (((phog-table2 pattern-table2) (extract-phog-stats corpus2)))
  (let ((grammar2 (learn-grammar corpus2 phog-table2 10 pattern-table2)))
    (test "grammar2 has let-stmt rule"
          #t
          (and (assq 'let-stmt grammar2) #t))

    (printf "\nLearned grammar2:\n")
    (print-grammar grammar2)

    ;; The let-stmt rule should parse both corpus examples
    (test "let-stmt parses 'let x = 42'"
          #t
          (parse? 'let-stmt "let x = 42"))

    (test "let-stmt parses 'let y = 7'"
          #t
          (parse? 'let-stmt "let y = 7"))

    ;; Verify grammar structurally
    (let ((results (verify-grammar corpus2)))
      (printf "\nVerification results (corpus2):\n")
      (for-each
        (lambda (r) (printf "  ~a: ~a\n" (car r) (cdr r)))
        results)

      (test "all corpus2 types verify ok"
            #t
            (let loop ((rest results))
              (or (null? rest)
                  (and (eq? (cdar rest) 'ok)
                       (loop (cdr rest)))))))))

;; ============================================================
;; Witten-Bell interpolation tests
;; ============================================================

(printf "\n=== Witten-Bell interpolation tests ===\n")

(printf "\n--- interpolate-distributions ---\n")

;; Pure MLE (lambda=1): backoff keys still present but with 0 weight
(let ((result (interpolate-distributions '((a . 1/2) (b . 1/2))
                                          '((a . 1/3) (c . 2/3))
                                          1)))
  (test "interpolate lambda=1 MLE keys correct"
        #t
        (and (= 1/2 (alist-ref 'a result 0))
             (= 1/2 (alist-ref 'b result 0))
             (= 0 (alist-ref 'c result 0)))))

;; Pure backoff (lambda=0): MLE keys still present but with 0 weight
(let ((result (interpolate-distributions '((a . 1/2) (b . 1/2))
                                          '((a . 1/3) (c . 2/3))
                                          0)))
  (test "interpolate lambda=0 backoff keys correct"
        #t
        (and (= 1/3 (alist-ref 'a result 0))
             (= 0 (alist-ref 'b result 0))
             (= 2/3 (alist-ref 'c result 0)))))

;; 50/50 blend
(let ((result (interpolate-distributions '((a . 1) (b . 0))
                                          '((a . 0) (b . 1))
                                          1/2)))
  (test "interpolate 50/50 blend"
        #t
        (and (= 1/2 (alist-ref 'a result 0))
             (= 1/2 (alist-ref 'b result 0)))))

;; Productions in backoff but not MLE get weight from backoff
(let ((result (interpolate-distributions '((a . 1))
                                          '((a . 1/2) (b . 1/2))
                                          1/2)))
  (test "unseen production gets backoff weight"
        #t
        (> (alist-ref 'b result 0) 0)))

(printf "\n--- phog-lookup-smooth ---\n")

;; Test with a hand-crafted nonterminal table
(let ((nt-table '(((expr 0 #f) . ((num . 10) (ident . 2)))
                   ((expr 0) . ((num . 8) (ident . 3) (str . 1)))
                   ((expr) . ((num . 20) (ident . 5) (str . 3) (bool . 1))))))
  ;; Full context should return results
  (let ((probs (phog-lookup-smooth nt-table '(expr 0 #f))))
    (test "phog-lookup-smooth returns results for known context"
          #t
          (not (null? probs)))
    (test "smoothed probs sum to ~1"
          #t
          (let ((total (apply + (map cdr probs))))
            (< (abs (- total 1)) 1/1000)))
    ;; bool should appear even though not at full context (from backoff)
    (test "backoff provides unseen production (bool)"
          #t
          (> (alist-ref 'bool probs 0) 0)))

  ;; Partial context should also work
  (let ((probs (phog-lookup-smooth nt-table '(expr 0))))
    (test "partial context returns results"
          #t
          (not (null? probs)))
    (test "partial smoothed probs sum to ~1"
          #t
          (let ((total (apply + (map cdr probs))))
            (< (abs (- total 1)) 1/1000))))

  ;; Completely unseen context should fall back
  (let ((probs (phog-lookup-smooth nt-table '(expr 99 foo))))
    (test "unseen context falls back to global"
          #t
          (not (null? probs)))
    (test "unseen context probs sum to ~1"
          #t
          (let ((total (apply + (map cdr probs))))
            (< (abs (- total 1)) 1/1000)))))

;; ============================================================
;; set primitive tests
;; ============================================================

(printf "\n=== set (byte-range matching) tests ===\n")

(printf "\n--- set constructor ---\n")

(test "set single range"
      '(set (48 . 57))
      (set '(48 . 57)))

(test "set multiple ranges"
      '(set (48 . 57) (65 . 90) (97 . 122))
      (set '(48 . 57) '(65 . 90) '(97 . 122)))

(test "set single integer (point)"
      '(set (42 . 42))
      (set 42))

(test "digit produces correct set node"
      '(set (48 . 57))
      (digit))

(test "lower produces correct set node"
      '(set (97 . 122))
      (lower))

(test "upper produces correct set node"
      '(set (65 . 90))
      (upper))

(test "alpha produces correct set node"
      '(set (65 . 90) (97 . 122))
      (alpha))

(test "alnum produces correct set node"
      '(set (48 . 57) (65 . 90) (97 . 122))
      (alnum))

(test "blank produces correct set node"
      '(set (9 . 10) (13 . 13) (32 . 32))
      (blank))

(printf "\n--- set matching (matcho) ---\n")

;; digit range matches '0' through '9'
(test "set digit matches #\\0"
      '(())
      (run* (rest) (matcho (digit) '(#\0) rest '())))

(test "set digit matches #\\5"
      '(())
      (run* (rest) (matcho (digit) '(#\5) rest '())))

(test "set digit matches #\\9"
      '(())
      (run* (rest) (matcho (digit) '(#\9) rest '())))

;; digit range rejects 'a'
(test "set digit rejects #\\a"
      '()
      (run* (rest) (matcho (digit) '(#\a) rest '())))

;; digit range rejects empty input
(test "set digit rejects empty"
      '()
      (run* (rest) (matcho (digit) '() rest '())))

;; lower matches letters
(test "set lower matches #\\m"
      '(())
      (run* (rest) (matcho (lower) '(#\m) rest '())))

(test "set lower rejects #\\M"
      '()
      (run* (rest) (matcho (lower) '(#\M) rest '())))

(printf "\n--- set with multiple ranges (alnum) ---\n")

;; alnum matches digits, upper, and lower
(test "set alnum matches #\\3"
      '(())
      (run* (rest) (matcho (alnum) '(#\3) rest '())))

(test "set alnum matches #\\Z"
      '(())
      (run* (rest) (matcho (alnum) '(#\Z) rest '())))

(test "set alnum matches #\\z"
      '(())
      (run* (rest) (matcho (alnum) '(#\z) rest '())))

(test "set alnum rejects #\\!"
      '()
      (run* (rest) (matcho (alnum) '(#\!) rest '())))

(printf "\n--- set in combination with cat/rep ---\n")

;; rep (digit) matches "123"
(test "rep digit matches \"123\""
      #t
      (and (member '() (run* (rest)
                          (matcho (rep (digit))
                                  (string->list "123")
                                  rest '())))
           #t))

;; cat of digit and lower matches "1a"
(test "cat digit lower matches \"1a\""
      '(())
      (run* (rest)
        (matcho (cat (digit) (lower))
                (string->list "1a")
                rest '())))

;; cat of digit and lower rejects "a1"
(test "cat digit lower rejects \"a1\""
      '()
      (run* (rest)
        (matcho (cat (digit) (lower))
                (string->list "a1")
                rest '())))

;; full-matcho with rep of digit
(test "full-matcho rep digit on \"42\""
      #t
      (not (null? (run 1 (q) (full-matcho (rep (digit))
                                           (string->list "42"))))))

;; set inside alt with lit
(test "alt of set and lit"
      '(())
      (run* (rest)
        (matcho (alt (digit) (lit #\-))
                '(#\-) rest '())))

(printf "\n--- set in parseo ---\n")

;; parseo with set produces correct output
(test "parseo set digit output"
      '((#\5))
      (run* (out) (full-parseo (digit) '(#\5) out (make-identity-reducer))))

(test "parseo rep digit output"
      #t
      (let ((results (run* (out) (full-parseo (rep (digit))
                                              (string->list "42")
                                              out
                                              (make-identity-reducer)))))
        ;; Should contain (#\4 #\2) among results
        (and (member '(#\4 #\2) results) #t)))

;; ============================================================
;; Anti-unification tests
;; ============================================================

(printf "\n=== Anti-unification tests ===\n")

(printf "\n--- helpers ---\n")

(test "iota 0" '() (iota 0))
(test "iota 3" '(0 1 2) (iota 3))

(test "transpose 2x3"
      '((a d) (b e) (c f))
      (transpose '((a b c) (d e f))))

(test "transpose 1x2"
      '((a) (b))
      (transpose '((a b))))

(test "group-by-skeleton single group"
      '(((a b c) (d e f)))
      (group-by-skeleton '((a b c) (d e f))))

(test "group-by-skeleton multiple groups by length"
      #t
      (let ((groups (group-by-skeleton '((a b c) (x) (d e f) (y)))))
        (and (= (length groups) 2)
             (= (length (car groups)) 2)
             (= (length (cadr groups)) 2))))

(test "group-by-skeleton splits same-length different-skeleton"
      2
      (length (group-by-skeleton
                '((*anonymous* pair *anonymous* pair *anonymous*)
                  (*anonymous* *anonymous* pair *anonymous* *anonymous*)
                  (*anonymous* pair *anonymous* string *anonymous*)))))

(printf "\n--- anti-unify-column ---\n")

;; Single constant value
(test "anti-unify-column single sym"
      '(sym num)
      (anti-unify-column 'expr '(num num num) 0 '() '((num plus num))))

;; Multiple distinct values produce alt
(test "anti-unify-column multiple syms"
      '(alt (sym plus) (sym minus))
      (anti-unify-column 'expr '(plus minus) 1 '() '((num plus num) (num minus num))))

(printf "\n--- anti-unify-group ---\n")

;; Same-length patterns: column-wise anti-unification
(test "anti-unify-group: (num plus num) (num minus num)"
      '(cat (sym num) (cat (alt (sym plus) (sym minus)) (sym num)))
      (anti-unify-group 'expr '((num plus num) (num minus num)) '()))

;; Single-element patterns
(test "anti-unify-group: single-element"
      '(sym num)
      (anti-unify-group 'expr '((num) (num)) '()))

(printf "\n--- anti-unify-patterns ---\n")

;; Mixed-length patterns: should produce top-level alt of length groups
(let ((result (anti-unify-patterns 'expr
                '((num plus num) (num minus num) (num))
                '()
                '())))
  (test "anti-unify-patterns produces alt"
        'alt
        (car result))
  ;; Should match all three patterns
  (test "anti-unify result matches (num plus num)"
        #t
        (not (null? (run 1 (rest) (sym-matcho result '(num plus num) rest '() #f 0 #f)))))
  (test "anti-unify result matches (num minus num)"
        #t
        (not (null? (run 1 (rest) (sym-matcho result '(num minus num) rest '() #f 0 #f)))))
  (test "anti-unify result matches (num)"
        #t
        (not (null? (run 1 (rest) (sym-matcho result '(num) rest '() #f 0 #f)))))
  ;; Should NOT match wrong patterns
  (test "anti-unify result rejects (plus num)"
        #t
        (null? (run 1 (rest) (sym-matcho result '(plus num) rest '() #f 0 #f)))))

;; Anti-unification with all-identical patterns produces simple cat
(let ((result (anti-unify-patterns 'stmt
                '((if expr stmt) (if expr stmt))
                '()
                '())))
  (test "identical patterns produce cat"
        '(cat (sym if) (cat (sym expr) (sym stmt)))
        result))

(printf "\n--- phog-order-regex ---\n")

;; Build a PHOG table where 'plus' is more frequent than 'minus'
(let ((phog-table '((expr . (((expr 1 #f #f num) . ((plus . 5) (minus . 2))))))))
  ;; alt with plus more frequent should put plus first
  (let ((ordered (phog-order-regex '(alt (sym minus) (sym plus))
                                    phog-table 'expr)))
    (test "phog-order puts more frequent branch first"
          '(alt (sym plus) (sym minus))
          ordered)))

;; Empty PHOG table: ordering should not crash, regex unchanged structurally
(let ((ordered (phog-order-regex '(alt (sym a) (sym b)) '() 'test)))
  (test "phog-order with empty table preserves branches"
        #t
        (and (pair? ordered) (eq? (car ordered) 'alt))))

;; Nested alt inside cat: both should be ordered
(let ((phog-table '((expr . (((expr 0 #f #f #f) . ((b . 10) (a . 1))))))))
  (let ((ordered (phog-order-regex '(cat (alt (sym a) (sym b)) (sym c))
                                    phog-table 'expr)))
    (test "phog-order recurses into cat children"
          '(cat (alt (sym b) (sym a)) (sym c))
          ordered)))

(printf "\n--- PHOG-weighted sym-matcho ---\n")

;; With PHOG weights, the alt branch should explore most-probable first
(let-values (((phog-table pattern-table) (extract-phog-stats corpus)))
  (let ((result (run 1 (rest)
                  (sym-matcho '(alt (sym minus) (sym plus))
                              '(plus)
                              rest
                              phog-table
                              'expr
                              0 #f))))
    (test "weighted sym-matcho finds match"
          '(())
          result))

  ;; Both branches still reachable
  (let ((result (run* (rest)
                  (sym-matcho '(alt (sym plus) (sym minus))
                              '(minus)
                              rest
                              phog-table
                              'expr
                              0 #f))))
    (test "weighted sym-matcho reaches second branch"
          '(())
          result)))

;; No PHOG data: should still work (graceful degradation)
(test "sym-matcho with empty PHOG table works"
      '(())
      (run* (rest) (sym-matcho '(alt (sym a) (sym b)) '(b) rest '() #f 0 #f)))

(printf "\n--- full-sym-matcho backward compat ---\n")

;; 2-arg form (backward compat)
(test "full-sym-matcho 2-arg form"
      #t
      (not (null? (run 1 (q) (full-sym-matcho '(sym num) '(num))))))

;; 4-arg form with PHOG data
(let-values (((phog-table pattern-table) (extract-phog-stats corpus)))
  (test "full-sym-matcho 4-arg form"
        #t
        (not (null? (run 1 (q)
                      (full-sym-matcho '(cat (sym num) (cat (alt (sym plus) (sym minus)) (sym num)))
                                       '(num plus num)
                                       phog-table 'expr))))))

;; ============================================================
;; End-to-end: learned grammar quality
;; ============================================================

(printf "\n=== End-to-end grammar quality ===\n")

(let-values (((phog-table pattern-table) (extract-phog-stats corpus)))
  (let ((grammar (learn-grammar corpus phog-table 10 pattern-table)))
    (let ((expr-rule (cdr (assq 'expr grammar))))
      ;; The expr rule should NOT contain _.0 (no unbound variables)
      (test "expr rule has no unbound vars"
            #f
            (let check ((r expr-rule))
              (cond
                ((and (symbol? r) (let ((s (symbol->string r)))
                                    (and (> (string-length s) 2)
                                         (char=? (string-ref s 0) #\_)
                                         (char=? (string-ref s 1) #\.))))
                 #t)
                ((pair? r) (or (check (car r)) (check (cdr r))))
                (else #f))))

      ;; The rule should be an alt of two branches (unary and binary)
      (test "expr rule is alt"
            'alt
            (car expr-rule))

      ;; The rule should parse all corpus examples
      (test "expr parses '43+2'" #t (parse? 'expr "43+2"))
      (test "expr parses '1-99'" #t (parse? 'expr "1-99"))
      (test "expr parses '7'" #t (parse? 'expr "7"))
      (test "expr parses '100+200'" #t (parse? 'expr "100+200"))

      ;; The rule should generalize to unseen but valid expressions
      (test "expr parses unseen '5+3'" #t (parse? 'expr "5+3"))
      (test "expr parses unseen '10-1'" #t (parse? 'expr "10-1")))))

;; ============================================================
;; mplus-biased tests
;; ============================================================

(printf "\n=== mplus-biased tests ===\n")

;; Test mplus-biased through miniKanren goals, which is how it's actually used.
;; k=1 should behave like fair mplus (all branches reachable)
(test "mplus-biased k=1 all results"
      4
      (length (run* (q)
                (lambda (st)
                  (let ((st^ (state-with-scope st (new-scope))))
                    (mplus-biased 1
                      (mplus ((== q 'a) st^) (lambda () ((== q 'b) st^)))
                      (lambda ()
                        (mplus ((== q 'x) st^) (lambda () ((== q 'y) st^))))))))))

;; k>1 should bias toward left stream — left results appear earlier
(let ((results (run 4 (q)
                 (lambda (st)
                   (let ((st^ (state-with-scope st (new-scope))))
                     (mplus-biased 3
                       (mplus* ((== q 'a) st^) ((== q 'b) st^) ((== q 'c) st^))
                       (lambda ()
                         (mplus ((== q 'x) st^) (lambda () ((== q 'y) st^))))))))))
  (test "mplus-biased k=3 first element from left"
        'a
        (car results))
  (test "mplus-biased k=3 second element from left"
        'b
        (cadr results))
  (test "mplus-biased k=3 third element from left"
        'c
        (caddr results)))

;; All elements reachable with k>1
(test "mplus-biased k=3 all 5 reachable"
      5
      (length (run* (q)
                (lambda (st)
                  (let ((st^ (state-with-scope st (new-scope))))
                    (mplus-biased 3
                      (mplus* ((== q 'a) st^) ((== q 'b) st^) ((== q 'c) st^))
                      (lambda ()
                        (mplus ((== q 'x) st^) (lambda () ((== q 'y) st^))))))))))

;; Empty left: should drain right
(test "mplus-biased empty left drains right"
      '(x)
      (run* (q)
        (lambda (st)
          (let ((st^ (state-with-scope st (new-scope))))
            (mplus-biased 5 #f (lambda () ((== q 'x) st^)))))))

;; Empty right: should drain left
(test "mplus-biased empty right drains left"
      '(a)
      (run* (q)
        (lambda (st)
          (let ((st^ (state-with-scope st (new-scope))))
            (mplus-biased 5
              ((== q 'a) st^)
              (lambda () #f))))))

;; ============================================================
;; condp-budgeted tests
;; ============================================================

(printf "\n=== condp-budgeted tests ===\n")

;; Highest weight should be explored first
(test "condp-budgeted explores highest weight first"
      '(42)
      (run 1 (q)
        (condp-budgeted
          1 (lambda (st) ((== q 'low) st))
          10 (lambda (st) ((== q 42) st))
          5 (lambda (st) ((== q 'mid) st)))))

;; All branches should be reachable
(test "condp-budgeted explores all branches"
      #t
      (let ((results (run* (q)
                       (condp-budgeted
                         1 (lambda (st) ((== q 'low) st))
                         10 (lambda (st) ((== q 'high) st))
                         5 (lambda (st) ((== q 'mid) st))))))
        (and (member 'low results)
             (member 'high results)
             (member 'mid results)
             #t)))

;; With large weight ratio, high-weight branch should dominate early results
(test "condp-budgeted bias in first N results"
      'high
      (car (run 1 (q)
             (condp-budgeted
               1 (lambda (st) ((== q 'low) st))
               100 (lambda (st) ((== q 'high) st))))))

;; Equal weights: all branches still reachable
(test "condp-budgeted equal weights all reachable"
      3
      (length (run* (q)
                (condp-budgeted
                  5 (lambda (st) ((== q 'a) st))
                  5 (lambda (st) ((== q 'b) st))
                  5 (lambda (st) ((== q 'c) st))))))

;; ============================================================
;; regex-weight-from-probs tests
;; ============================================================

(printf "\n=== regex-weight-from-probs tests ===\n")

(test "regex-weight-from-probs sym found"
      3/10
      (regex-weight-from-probs '(sym num) '((num . 3/10) (ident . 7/10))))

(test "regex-weight-from-probs sym not found"
      1/1000
      (regex-weight-from-probs '(sym unknown) '((num . 3/10) (ident . 7/10))))

(test "regex-weight-from-probs cat uses first child"
      3/10
      (regex-weight-from-probs '(cat (sym num) (sym plus)) '((num . 3/10) (plus . 7/10))))

(test "regex-weight-from-probs empty probs"
      1
      (regex-weight-from-probs '(sym num) '()))

(test "regex-weight-from-probs non-sym defaults to 1"
      1
      (regex-weight-from-probs '(rep (sym num)) '((num . 5/10))))

;; ============================================================
;; Context-specific sym-matcho tests
;; ============================================================

(printf "\n=== Context-specific sym-matcho tests ===\n")

;; Build a PHOG table with position-specific counts where:
;; At position 0, 'num' is very frequent (weight 20)
;; At position 1, 'plus' is frequent (weight 10), 'minus' rare (weight 1)
(let ((ctx-phog '((expr . (((expr 0 expr #f #f) . ((num . 20) (ident . 1)))
                            ((expr 1 expr #f num) . ((plus . 10) (minus . 1))))))))

  ;; Context-specific weighting: at position 1 with lsib=num, plus should be preferred
  (let ((result (run 2 (rest)
                  (sym-matcho '(alt (sym minus) (sym plus))
                              '(plus)
                              rest
                              ctx-phog
                              'expr
                              1 'num))))
    (test "context-specific sym-matcho finds match with position context"
          '(())
          result))

  ;; Both branches still reachable with context
  (let ((result (run* (rest)
                  (sym-matcho '(alt (sym plus) (sym minus))
                              '(minus)
                              rest
                              ctx-phog
                              'expr
                              1 'num))))
    (test "context-specific sym-matcho second branch reachable"
          '(())
          result)))

;; ============================================================
;; Position threading tests
;; ============================================================

(printf "\n=== Position threading tests ===\n")

;; After matching cat of (sym a) (sym b), position should advance
;; We test indirectly: matching works correctly with position tracking
(test "cat position threading basic"
      '(())
      (run* (rest) (sym-matcho '(cat (sym a) (sym b)) '(a b) rest '() #f 0 #f)))

;; Deeper cat chains
(test "cat position threading 3 elements"
      '(())
      (run* (rest) (sym-matcho '(cat (sym a) (cat (sym b) (sym c)))
                                '(a b c) rest '() #f 0 #f)))

;; Position threading with alt inside cat
(test "position threading with alt in cat"
      #t
      (let ((regex '(cat (sym num) (cat (alt (sym plus) (sym minus)) (sym num)))))
        (and (not (null? (run 1 (rest) (sym-matcho regex '(num plus num) rest '() #f 0 #f))))
             (not (null? (run 1 (rest) (sym-matcho regex '(num minus num) rest '() #f 0 #f)))))))

;; Position threading with PHOG table: verify that context-at-position is used
;; Build a PHOG table where position 2 has different weights than position 0
(let ((pos-phog '((stmt . (((stmt 0 stmt #f #f) . ((if . 5) (while . 1)))
                            ((stmt 2 stmt #f expr) . ((stmt . 10) (expr . 2))))))))
  ;; At pos=2 with lsib=expr, 'stmt' should be weighted higher than 'expr'
  ;; but both should match
  (test "position-aware PHOG at pos=2"
        '(())
        (run 1 (rest)
          (sym-matcho '(alt (sym expr) (sym stmt))
                      '(stmt)
                      rest
                      pos-phog
                      'stmt
                      2 'expr))))

;; rep position threading: full consumption via full-sym-matcho
(test "rep position threading"
      #t
      (not (null? (run 1 (q) (full-sym-matcho '(rep (sym a)) '(a a a))))))

;; opt position threading
(test "opt position threading (match)"
      '(())
      (run 1 (rest)
        (sym-matcho '(cat (opt (sym a)) (sym b)) '(a b) rest '() #f 0 #f)))

(test "opt position threading (skip)"
      '(())
      (run 1 (rest)
        (sym-matcho '(cat (opt (sym a)) (sym b)) '(b) rest '() #f 0 #f)))

;; ============================================================
;; Summary
;; ============================================================

(test-summary)