cleaning sesh

src/core/batch.rkt

@@ -5,7 +5,7 @@
          "../syntax/types.rkt")
 
 (provide progs->batch ; (Listof Expr) -> Batch
-         batch->progs ; Batch -> (Listof Expr)
+         batch->progs ; Batch -> *(or (Listof Root) (Vectorof Root)) -> (Listof Expr)
          (struct-out batch)
          (struct-out batchref) ; temporarily for patch.rkt
          (struct-out mutable-batch) ; temporarily for patch.rkt
@@ -15,8 +15,7 @@
          batch-replace ; Batch -> (Expr<Batchref> -> Expr<Batchref>) -> Batch
          egg-nodes->batch ; Nodes -> Spec-maps -> Batch -> (Listof Batchref)
          batchref->expr ; Batchref -> Expr
-         batch-extract-exprs ; Batch -> (Listof Root) -> (Listof Expr)
-         remove-zombie-nodes ; Batch -> Batch
+         batch-remove-zombie ; Batch -> *(Vectorof Root) -> Batch
          mutable-batch-add-expr! ; Mutable-batch -> Root
          mutable-batch->batch ; Mutable-batch -> Batch
          batch->mutable-batch ; Batch -> Mutable-batch
@@ -95,22 +94,14 @@
     (batch-push! b (expr-recurse prog munge)))
   (munge expr))
 
-(define (batch-extract-exprs b roots)
+(define (batch->progs b [roots (batch-roots b)])
   (define exprs (make-vector (batch-length b)))
   (for ([node (in-vector (batch-nodes b))]
         [idx (in-naturals)])
     (vector-set! exprs idx (expr-recurse node (lambda (x) (vector-ref exprs x)))))
   (for/list ([root roots])
     (vector-ref exprs root)))
 
-(define (batch->progs b)
-  (define exprs (make-vector (batch-length b)))
-  (for ([node (in-vector (batch-nodes b))]
-        [idx (in-naturals)])
-    (vector-set! exprs idx (expr-recurse node (lambda (x) (vector-ref exprs x)))))
-  (for/list ([root (batch-roots b)])
-    (vector-ref exprs root)))
-
 (define (batch-replace b f)
   (define out (make-mutable-batch))
   (define mapping (make-vector (batch-length b) -1))
@@ -134,9 +125,8 @@
 ; The function removes any zombie nodes from batch with respect to the roots
 ; Time complexity: O(|R| + |N|), where |R| - number of roots, |N| - length of nodes
 ; Space complexity: O(|N| + |N*| + |R|), where |N*| is a length of nodes without zombie nodes
-(define (remove-zombie-nodes input-batch)
+(define (batch-remove-zombie input-batch [roots (batch-roots input-batch)])
   (define nodes (batch-nodes input-batch))
-  (define roots (batch-roots input-batch))
   (define nodes-length (batch-length input-batch))
 
   (define zombie-mask (make-vector nodes-length #t))
@@ -271,7 +261,7 @@
   (require rackunit)
   (define (zombie-test #:nodes nodes #:roots roots)
     (define in-batch (batch nodes roots '()))
-    (define out-batch (remove-zombie-nodes in-batch))
+    (define out-batch (batch-remove-zombie in-batch))
     (check-equal? (batch->progs out-batch) (batch->progs in-batch))
     (batch-nodes out-batch))
 

src/core/egg-herbie.rkt

@@ -135,7 +135,7 @@
   (set-batch-roots! batch roots) ; make sure that we work with the right roots
   ; the algorithm may crash if batch-length is zero
   (define insert-batch
-    (if (zero? (batch-length batch)) batch (remove-zombie-nodes (batch-parse-approx batch))))
+    (if (zero? (batch-length batch)) batch (batch-remove-zombie (batch-parse-approx batch))))
 
   (define mappings (build-vector (batch-length insert-batch) values))
   (define (remap x)

src/core/localize.rkt

@@ -49,13 +49,16 @@
 
   ; run egg
   (define simplified
-    (simplify-batch runner
-                    (typed-egg-batch-extractor
-                     (if (*egraph-platform-cost*) platform-egg-cost-proc default-egg-cost-proc)
-                     batch)))
+    (map batchref->expr
+         (map last
+              (simplify-batch runner
+                              (typed-egg-batch-extractor (if (*egraph-platform-cost*)
+                                                             platform-egg-cost-proc
+                                                             default-egg-cost-proc)
+                                                         batch)))))
 
   ; run egg
-  (define simplifiedss (regroup-nested subexprss (map batchref->expr (map last simplified))))
+  (define simplifiedss (regroup-nested subexprss simplified))
 
   ; build map from starting expr to simplest
   (define expr->simplest (make-hash))

src/core/preprocess.rkt

@@ -80,9 +80,10 @@
   ; alternatives
   (define start-alt (make-alt expr))
   (cons start-alt
-        (remove-duplicates (for/list ([expr (rest simplified)])
-                             (alt expr `(simplify () ,runner #f #f) (list start-alt) '()))
-                           alt-equal?)))
+        (remove-duplicates
+         (for/list ([expr (rest simplified)])
+           (alt (batchref->expr expr) `(simplify () ,runner #f #f) (list start-alt) '()))
+         alt-equal?)))
 
 ;; See https://pavpanchekha.com/blog/symmetric-expressions.html
 (define (find-preprocessing init expr ctx)

src/core/programs.rkt

@@ -41,7 +41,7 @@
     [(? literal?) (get-representation (literal-precision node))]
     [(? variable?) (context-lookup ctx node)]
     [(list '$approx _ impl)
-     (repr-of-node batch impl ctx)] ; here is a hack to match an after-parse structure
+     (repr-of-node batch impl ctx)] ; this is ugly, but egraph-add-exprs relies on this match
     [(list 'if cond ift iff) (repr-of-node batch ift ctx)]
     [(list op args ...) (impl-info op 'otype)]))
 

src/core/simplify.rkt

@@ -21,8 +21,7 @@
 ;; if the input specifies proofs, it instead returns proofs for these expressions
 (define/contract (simplify-batch runner extractor)
   (-> egg-runner? procedure? (listof (listof batchref?)))
-  (timeline-push! 'inputs
-                  (map ~a (batch-extract-exprs (egg-runner-batch runner) (egg-runner-roots runner))))
+  (timeline-push! 'inputs (map ~a (batch->progs (egg-runner-batch runner) (egg-runner-roots runner))))
   (timeline-push! 'method "egg-herbie")
 
   (define simplifieds (run-egg runner (cons 'single extractor)))

src/core/taylor.rkt

@@ -26,16 +26,14 @@
     (define (next [iter 0])
       (define coeff (simplify (replace-expression (coeffs i) var ((cdr tform) var))))
       (set! i (+ i 1))
-      (define out
-        (match coeff
-          [0
-           (if (< iter iters)
-               (next (+ iter 1))
-               (simplify (make-horner ((cdr tform) var) (reverse terms))))]
-          [_
-           (set! terms (cons (cons coeff (- i offset 1)) terms))
-           (simplify (make-horner ((cdr tform) var) (reverse terms)))]))
-      out)
+      (match coeff
+        [0
+         (if (< iter iters)
+             (next (+ iter 1))
+             (simplify (make-horner ((cdr tform) var) (reverse terms))))]
+        [_
+         (set! terms (cons (cons coeff (- i offset 1)) terms))
+         (simplify (make-horner ((cdr tform) var) (reverse terms)))]))
     next))
 
 ;; Our Taylor expander prefers sin, cos, exp, log, neg over trig, htrig, pow, and subtraction