code simplification. Particularly at simplify.rkt

src/core/localize.rkt

@@ -49,7 +49,7 @@
 
   ; run egg
   (define simplified
-    (map (compose debatchref last)
+    (map debatchref
          (simplify-batch runner
                          (typed-egg-batch-extractor
                           (if (*egraph-platform-cost*) platform-egg-cost-proc default-egg-cost-proc)

src/core/mainloop.rkt

@@ -377,7 +377,7 @@
 
      ; run egg
      (define simplified
-       (map (compose debatchref last)
+       (map debatchref
             (simplify-batch runner
                             (typed-egg-batch-extractor (if (*egraph-platform-cost*)
                                                            platform-egg-cost-proc

src/core/patch.rkt

@@ -59,16 +59,14 @@
   (define simplified
     (reap [sow]
           (for ([altn (in-list approxs)]
-                [outputs (in-list simplification-options)])
-            (match-define (cons _ simplified) outputs)
+                [batchreff (in-list simplification-options)])
             (define prev (car (alt-prevs altn)))
-            (for ([expr (in-list simplified)])
-              (define spec (prog->spec (debatchref (alt-expr prev))))
-              (define idx
-                (batch-push! global-batch-mutable
-                             (approx (mutable-batch-munge! global-batch-mutable spec)
-                                     (batchref-idx expr))))
-              (sow (alt (batchref global-batch idx) `(simplify ,runner #f #f) (list altn) '()))))))
+            (define spec (prog->spec (debatchref (alt-expr prev))))
+            (define idx
+              (batch-push! global-batch-mutable
+                           (approx (mutable-batch-munge! global-batch-mutable spec)
+                                   (batchref-idx batchreff))))
+            (sow (alt (batchref global-batch idx) `(simplify ,runner #f #f) (list altn) '())))))
 
   ; Commit changes to global-batch
   (batch-copy-mutable-nodes! global-batch global-batch-mutable)
@@ -188,12 +186,11 @@
   ; Batch to where we will extract everything
   ; Roots of this batch are constantly updated
   (define global-batch (progs->batch exprs))
-  (define start-roots (batch-roots global-batch))
 
   ; Starting alternatives
   (define start-altns
     (for/list ([expr (in-list exprs)]
-               [root (in-vector start-roots)])
+               [root (batch-roots global-batch)])
       (define repr (repr-of expr (*context*)))
       (alt (batchref global-batch root) (list 'patch expr repr) '() '())))
 
@@ -202,6 +199,4 @@
   ; Recursive rewrite
   (define rewritten (if (flag-set? 'generate 'rr) (run-rr start-altns global-batch) '()))
 
-  (define out (append approximations rewritten))
-
-  out)
+  (append approximations rewritten))

src/core/preprocess.rkt

@@ -72,18 +72,18 @@
 
   ; 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 debatchref
+         (simplify-batch runner
+                         (typed-egg-batch-extractor
+                          (if (*egraph-platform-cost*) platform-egg-cost-proc default-egg-cost-proc)
+                          batch))))
 
   ; alternatives
   (define start-alt (make-alt expr))
   (cons start-alt
-        (remove-duplicates
-         (for/list ([expr (rest simplified)])
-           (alt (debatchref expr) `(simplify () ,runner #f #f) (list start-alt) '()))
-         alt-equal?)))
+        (remove-duplicates (for/list ([expr simplified])
+                             (alt 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/simplify.rkt

@@ -20,20 +20,12 @@
 ;; the last expression is the simplest unless something went wrong due to unsoundness
 ;; if the input specifies proofs, it instead returns proofs for these expressions
 (define/contract (simplify-batch runner extractor)
-  (-> egg-runner? procedure? (listof (listof batchref?)))
+  (-> egg-runner? procedure? (listof batchref?))
   (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)))
-  (define out
-    (for/list ([simplified simplifieds]
-               [root (egg-runner-roots runner)])
-      (if (equal? root (batchref-idx (car simplified)))
-          simplified
-          (cons (batchref (egg-runner-batch runner) root) simplified))))
-
-  (timeline-push! 'outputs (map ~a (apply append out)))
-  out)
+  (define simplifieds (apply append (run-egg runner (cons 'single extractor))))
+  (timeline-push! 'outputs (map (compose ~a debatchref) simplifieds))
+  simplifieds)
 
 (module+ test
   (require "../syntax/types.rkt"
@@ -58,7 +50,7 @@
                        (map (lambda (_) 'real) args)
                        `((,(*simplify-rules*) . ((node . ,(*node-limit*)))))))
     (define extractor (typed-egg-batch-extractor default-egg-cost-proc batch))
-    (map (compose debatchref last) (simplify-batch runner extractor)))
+    (map debatchref (simplify-batch runner extractor)))
 
   (define test-exprs
     '((1 . 1) (0 . 0)