From 562cecd10ca94f432069886a5da8d69f3fefbf0e Mon Sep 17 00:00:00 2001
From: Max New <maxsnew@gmail.com>
Date: Fri, 8 Nov 2013 17:40:09 -0600
Subject: [PATCH] Add disjoint sum enum and fix enum

---
 .../redex-lib/redex/private/enum.rkt          |  27 +++--
 .../redex-lib/redex/private/enumerator.rkt    | 106 ++++++++++++++----
 .../redex/tests/enumerator-test.rkt           |  52 +++++++++
 3 files changed, 156 insertions(+), 29 deletions(-)

diff --git a/pkgs/redex-pkgs/redex-lib/redex/private/enum.rkt b/pkgs/redex-pkgs/redex-lib/redex/private/enum.rkt
index 6f4a7a3d50..d144f642db 100644
--- a/pkgs/redex-pkgs/redex-lib/redex/private/enum.rkt
+++ b/pkgs/redex-pkgs/redex-lib/redex/private/enum.rkt
@@ -87,7 +87,7 @@
   (define (loop pat)
     (match-a-pattern
      pat
-     [`any (sum/e any/e (many/e any/e))]
+     [`any any/e]
      [`number num/e]
      [`string string/e]
      [`natural natural/e]
@@ -175,6 +175,7 @@
      (define term
        (fill-refs refpat))
      (λ (_) term)]
+    
     [(name-ref n)
      (λ (nv)
         (t-env-name-ref nv n))]
@@ -218,10 +219,11 @@
    (many/e char/e)))
 
 (define integer/e
-  (sum/e nats
-         (map/e (λ (n) (- (+ n 1)))
-                (λ (n) (- (- n) 1))
-                nats)))
+  (disj-sum/e (cons nats (λ (n) (>= n 0)))
+              (cons (map/e (λ (n) (- (+ n 1)))
+                           (λ (n) (- (- n) 1))
+                           nats)
+                    (λ (n) (< n 0)))))
 
 ;; This is really annoying so I turned it off
 (define real/e empty/e)
@@ -238,11 +240,18 @@
    (compose string->list symbol->string)
    (many1/e char/e)))
 
+(define base/e
+  (disj-sum/e (cons (const/e '()) null?)
+              (cons num/e number?)
+              (cons string/e string?)
+              (cons bool/e boolean?)
+              (cons var/e symbol?)))
+
 (define any/e
-  (sum/e num/e
-         string/e
-         bool/e
-         var/e))
+  (fix/e +inf.f
+         (λ (any/e)
+            (disj-sum/e (cons base/e (negate pair?))
+                        (cons (cons/e any/e any/e) pair?)))))
 
 (define (unsupported pat)
   (error 'generate-term "#:i-th does not support ~s patterns" pat))
diff --git a/pkgs/redex-pkgs/redex-lib/redex/private/enumerator.rkt b/pkgs/redex-pkgs/redex-lib/redex/private/enumerator.rkt
index a4fc88e1d4..332bf19f4a 100644
--- a/pkgs/redex-pkgs/redex-lib/redex/private/enumerator.rkt
+++ b/pkgs/redex-pkgs/redex-lib/redex/private/enumerator.rkt
@@ -3,6 +3,7 @@
          racket/match
          racket/list
          racket/function
+         racket/promise
          data/gvector)
 
 (provide enum
@@ -14,6 +15,7 @@
          const/e
          from-list/e
          sum/e
+         disj-sum/e
          cons/e
          dep/e
          dep2/e ;; requires size (eventually should replace dep/e with this)
@@ -21,6 +23,7 @@
          filter/e ;; very bad, only use for small enums
          except/e 
          thunk/e
+         fix/e
          many/e
          many1/e
          list/e
@@ -176,14 +179,17 @@
 ;; input list should not contain duplicates
 ;; equality on eq?
 (define (from-list/e l)
+  (define rev-map
+    (for/hash ([i (in-naturals)]
+               [x (in-list l)])
+      (values x i)))
   (if (empty? l)
       empty/e
       (enum (length l)
             (λ (n)
                (list-ref l n))
             (λ (x)
-               (length (take-while l (λ (y)
-                                        (not (eq? x y)))))))))
+               (hash-ref rev-map x)))))
 
 ;; take-while : Listof a, (a -> bool) -> Listof a
 (define (take-while l pred)
@@ -278,6 +284,63 @@
                  (map-pairs/even (cdr l)))))
      (apply sum/e (map-pairs sum/e identity (list* a b c rest)))]))
 
+(define (disj-sum/e e-p . e-ps)
+  (define/match (disj-sum2/e e-p1 e-p2)
+    [((cons e1 1?) (cons e2 2?))
+     ;; Sum two enumerators of different sizes
+     (define (sum-uneven less/e less? more/e more?)
+       ;; interleave until less/e is exhausted
+       ;; pairsdone is 1+ the highest index using less/e
+       (define less-size (size less/e))
+       (define pairsdone (* 2 less-size))
+       (define (from-nat n)
+         (if (< n pairsdone)
+             (let-values ([(q r) (quotient/remainder n 2)])
+               ;; Always put e1 first though!
+               (decode (match r
+                         [0 e1]
+                         [1 e2])
+                       q))
+             (decode more/e (- n less-size))))
+       (define (to-nat x)
+         (cond [(less? x)
+                (* 2 (encode less/e x))]
+               [(more? x)
+                (define i (encode more/e x))
+                (if (< i less-size)
+                    (+ (* 2 i) 1)
+                    (+ (- i less-size) pairsdone))]
+               [else (error "bad term")]))
+       (enum (+ less-size (size more/e))
+             from-nat
+             to-nat))
+     (define s1 (size e1))
+     (define s2 (size e2))
+     (cond [(= 0 s1) e2]
+           [(= 0 s2) e1]
+           [(< s1 s2) (sum-uneven e1 1? e2 2?)]
+           [(< s2 s1) (sum-uneven e2 2? e1 1?)]
+           [else ;; both the same length, interleave them
+            (define (from-nats n)
+              (cond [(even? n) (decode e1 (/ n 2))]
+                    [else (decode e2 (/ (- n 1) 2))]))
+            (define (to-nats x)
+              (cond [(1? x) (* (encode e1 x) 2)]
+                    [(2? x) (+ 1 (* (encode e2 x) 2))]
+                    [else (error "bad term")]))
+            (enum (+ s1 s2) from-nats to-nats)])])
+  (car
+   (foldr (λ (e-p1 e-p2)
+             (match* (e-p1 e-p2)
+                     [((cons e1 1?) (cons e2 2?))
+                      (cons (disj-sum2/e e-p1
+                                         (cons e2 (negate 1?)))
+                            (λ (x)
+                               (or (1? x)
+                                   (2? x))))]))
+          (cons empty/e (λ (_) #f))
+          (cons e-p e-ps))))
+
 (define n*n
   (enum +inf.f
         (λ (n)
@@ -695,30 +758,33 @@
 
 ;; thunk/e : Nat or +-Inf, ( -> enum a) -> enum a
 (define (thunk/e s thunk)
-  (let* ([e #f]
-         [get-e (λ ()
-                   (or e
-                       (and (set! e (thunk))
-                            e)))])
-    (enum s
-          (λ (n)
-             (decode (get-e) n))
-          (λ (x)
-             (encode (get-e) x)))))
+  (define promise/e (delay (thunk)))
+  (enum s
+        (λ (n)
+           (decode (force promise/e) n))
+        (λ (x)
+           (encode (force promise/e) x))))
+
+;; fix/e : size (enum a -> enum a) -> enum a
+(define (fix/e size f/e)
+  (define self (delay (f/e (fix/e size f/e))))
+  (enum size
+        (λ (n)
+           (decode (force self) n))
+        (λ (x)
+           (encode (force self) x))))
 
 ;; many/e : enum a -> enum (listof a)
 ;;       or : enum a, #:length natural -> enum (listof a)
 (define many/e
   (case-lambda
     [(e)
-     (thunk/e
-      (if (= 0 (size e))
-          0
-          +inf.f)
-      (λ ()
-         (sum/e
-          (const/e '())
-          (cons/e e (many/e e)))))]
+     (fix/e (if (= 0 (size e))
+                0
+                +inf.f)
+            (λ (self)
+               (disj-sum/e (cons (const/e '()) null?)
+                           (cons (cons/e e self) pair?))))]
     [(e n)
      (list/e (build-list n (const e)))]))
 
diff --git a/pkgs/redex-pkgs/redex-test/redex/tests/enumerator-test.rkt b/pkgs/redex-pkgs/redex-test/redex/tests/enumerator-test.rkt
index 1cf66953fd..531c4924bc 100644
--- a/pkgs/redex-pkgs/redex-test/redex/tests/enumerator-test.rkt
+++ b/pkgs/redex-pkgs/redex-test/redex/tests/enumerator-test.rkt
@@ -125,6 +125,53 @@
    (check-equal? (encode odd-or-even 3) 3)
    (check-bijection? odd-or-even)))
 
+(test-begin
+ (define bool-or-num
+   (disj-sum/e (cons bools/e boolean?)
+               (cons (from-list/e '(0 1 2 3)) number?)))
+ (define bool-or-nat
+   (disj-sum/e (cons bools/e boolean?)
+              (cons nats number?)))
+ (define nat-or-bool
+   (disj-sum/e (cons nats number?)
+              (cons bools/e boolean?)))
+ (define odd-or-even
+   (disj-sum/e (cons evens/e even?)
+          (cons odds/e odd?)))
+ (check-equal? (size bool-or-num) 6)
+   
+ (check-equal? (decode bool-or-num 0) #t)
+ (check-equal? (decode bool-or-num 1) 0)
+ (check-equal? (decode bool-or-num 2) #f)
+ (check-equal? (decode bool-or-num 3) 1)
+ (check-equal? (decode bool-or-num 4) 2)
+ (check-equal? (decode bool-or-num 5) 3)
+   
+ (check-exn exn:fail?
+            (λ ()
+               (decode bool-or-num 6)))
+ (check-bijection? bool-or-num)
+   
+ (check-equal? (size bool-or-nat)
+               +inf.f)
+ (check-equal? (decode bool-or-nat 0) #t)
+ (check-equal? (decode bool-or-nat 1) 0)
+ (check-bijection? bool-or-nat)
+   
+ (check-equal? (size odd-or-even)
+               +inf.f)
+ (check-equal? (decode odd-or-even 0) 0)
+ (check-equal? (decode odd-or-even 1) 1)
+ (check-equal? (decode odd-or-even 2) 2)
+ (check-exn exn:fail?
+            (λ ()
+               (decode odd-or-even -1)))
+ (check-equal? (encode odd-or-even 0) 0)   
+ (check-equal? (encode odd-or-even 1) 1)
+ (check-equal? (encode odd-or-even 2) 2)
+ (check-equal? (encode odd-or-even 3) 3)
+ (check-bijection? odd-or-even))
+
 ;; cons/e tests
 (define bool*bool (cons/e bools/e bools/e))
 (define 1*b (cons/e (const/e 1) bools/e))
@@ -354,3 +401,8 @@
       (except/e (up-to n) excepts))
    '(2 4 6)))
 (check-bijection? complicated)
+
+;; many/e tests
+(define natss
+  (many/e nats))
+(check-bijection? natss)