Merge pull request #1330 from herbie-fp/no-immutable-batches

No immutable batches

Author: pavpanchekha

src/core/alt-table.rkt

@@ -28,8 +28,8 @@
 
 (define (alt-batch-cost batch repr)
   (define node-cost-proc (platform-node-cost-proc (*active-platform*)))
-  (define costs (make-vector (vector-length (batch-nodes batch)) 0))
-  (for ([node (in-vector (batch-nodes batch))]
+  (define costs (make-vector (batch-length batch) 0))
+  (for ([node (in-batch batch)]
         [i (in-naturals)])
     (define cost
       (match node

src/core/batch.rkt

@@ -2,38 +2,45 @@
 
 (require "../syntax/syntax.rkt"
          "../utils/common.rkt"
-         "../utils/alternative.rkt") ; for unbatchify-alts
+         "../utils/alternative.rkt" ; for unbatchify-alts
+         "dvector.rkt")
 
 (provide progs->batch ; List<Expr> -> Batch
          batch->progs ; Batch -> ?(or List<Root> Vector<Root>) -> List<Expr>
+
          (struct-out batch)
-         (struct-out batchref)
-         (struct-out mutable-batch)
+         make-batch ; Batch
+         batch-push! ; Batch -> Node -> Idx
+         batch-munge! ; Batch -> Expr -> Root
+         batch-copy ; Batch -> Batch
          batch-length ; Batch -> Integer
          batch-tree-size ; Batch -> Integer
          batch-free-vars
-         batch-ref ; Batch -> Idx -> Expr
-         deref ; Batchref -> Expr
+         in-batch ; Batch -> Sequence<Node>
+         batch-ref ; Batch -> Idx -> Node
+         batch-pull ; Batch -> Idx -> Expr
          batch-replace ; Batch -> (Expr<Batchref> -> Expr<Batchref>) -> Batch
-         debatchref ; Batchref -> Expr
          batch-alive-nodes ; Batch -> ?Vector<Root> -> Vector<Idx>
          batch-reconstruct-exprs ; Batch -> Vector<Expr>
          batch-remove-zombie ; Batch -> ?Vector<Root> -> Batch
-         mutable-batch-munge! ; Mutable-batch -> Expr -> Root
-         make-mutable-batch ; Mutable-batch
-         batch->mutable-batch ; Batch -> Mutable-batch
-         batch-copy-mutable-nodes! ; Batch -> Mutable-batch -> Void
-         mutable-batch-push! ; Mutable-batch -> Node -> Idx
-         batch-copy
+
+         (struct-out batchref)
+         deref ; Batchref -> Expr
+         debatchref ; Batchref -> Expr
+
          unbatchify-alts)
 
 ;; Batches store these recursive structures, flattened
-(struct batch ([nodes #:mutable] [roots #:mutable]))
-
-(struct mutable-batch ([nodes #:mutable] [index #:mutable] cache))
+(struct batch ([nodes #:mutable] [index #:mutable] cache [roots #:mutable]))
 
 (struct batchref (batch idx) #:transparent)
 
+(define (make-batch)
+  (batch (make-dvector) (make-hash) (make-hasheq) (vector)))
+
+(define (in-batch batch [start 0] [end #f] [step 1])
+  (in-dvector (batch-nodes batch) start end step))
+
 ;; This function defines the recursive structure of expressions
 (define (expr-recurse expr f)
   (match expr
@@ -55,69 +62,58 @@
   (map (curry alt-map unmunge) altns))
 
 (define (batch-length b)
-  (cond
-    [(batch? b) (vector-length (batch-nodes b))]
-    [(mutable-batch? b) (hash-count (mutable-batch-index b))]
-    [else (error 'batch-length "Invalid batch" b)]))
-
-(define (make-mutable-batch)
-  (mutable-batch '() (make-hash) (make-hasheq)))
+  (dvector-length (batch-nodes b)))
 
-(define (mutable-batch-push! b term)
-  (define hashcons (mutable-batch-index b))
+(define (batch-push! b term)
+  (define hashcons (batch-index b))
   (hash-ref! hashcons
              term
              (lambda ()
-               (define new-idx (hash-count hashcons))
-               (hash-set! hashcons term new-idx)
-               (set-mutable-batch-nodes! b (cons term (mutable-batch-nodes b)))
-               new-idx)))
-
-(define (mutable-batch->batch b roots)
-  (batch (list->vector (reverse (mutable-batch-nodes b))) roots))
-
-(define (batch->mutable-batch b)
-  (mutable-batch (reverse (vector->list (batch-nodes b))) (batch-restore-index b) (make-hasheq)))
-
-(define (batch-copy-mutable-nodes! b mb)
-  (set-batch-nodes! b (list->vector (reverse (mutable-batch-nodes mb)))))
+               (define idx (hash-count hashcons))
+               (hash-set! hashcons term idx)
+               (dvector-add! (batch-nodes b) term)
+               idx)))
 
 (define (batch-copy b)
-  (batch (vector-copy (batch-nodes b)) (vector-copy (batch-roots b))))
+  (batch (dvector-copy (batch-nodes b))
+         (hash-copy (batch-index b))
+         (hash-copy (batch-cache b))
+         (vector-copy (batch-roots b))))
 
 (define (deref x)
   (match-define (batchref b idx) x)
-  (expr-recurse (vector-ref (batch-nodes b) idx) (lambda (ref) (batchref b ref))))
+  (expr-recurse (batch-ref b idx) (lambda (ref) (batchref b ref))))
 
 (define (debatchref x)
   (match-define (batchref b idx) x)
-  (batch-ref b idx))
+  (batch-pull b idx))
 
 (define (progs->batch exprs #:vars [vars '()])
-  (define out (make-mutable-batch))
+  (define out (make-batch))
 
   (for ([var (in-list vars)])
-    (mutable-batch-push! out var))
+    (batch-push! out var))
   (define roots
     (for/vector #:length (length exprs)
                 ([expr (in-list exprs)])
-      (mutable-batch-munge! out expr)))
+      (batch-munge! out expr)))
 
-  (mutable-batch->batch out roots))
+  (set-batch-roots! out roots)
+  out)
 
 (define (batch-tree-size b)
-  (define len (vector-length (batch-nodes b)))
+  (define len (batch-length b))
   (define counts (make-vector len 0))
   (for ([i (in-naturals)]
-        [node (in-vector (batch-nodes b))])
+        [node (in-batch b)])
     (define args (reap [sow] (expr-recurse node sow)))
     (vector-set! counts i (apply + 1 (map (curry vector-ref counts) args))))
   (apply + (map (curry vector-ref counts) (vector->list (batch-roots b)))))
 
-(define (mutable-batch-munge! b expr)
-  (define cache (mutable-batch-cache b))
+(define (batch-munge! b expr)
+  (define cache (batch-cache b))
   (define (munge prog)
-    (hash-ref! cache prog (lambda () (mutable-batch-push! b (expr-recurse prog munge)))))
+    (hash-ref! cache prog (lambda () (batch-push! b (expr-recurse prog munge)))))
   (munge expr))
 
 (define (batch->progs b [roots (batch-roots b)])
@@ -126,9 +122,9 @@
     (vector-ref exprs root)))
 
 (define (batch-free-vars batch)
-  (define out (make-vector (vector-length (batch-nodes batch))))
+  (define out (make-vector (batch-length batch)))
   (for ([i (in-naturals)]
-        [node (in-vector (batch-nodes batch))])
+        [node (in-batch batch)])
     (define fv
       (cond
         [(symbol? node) (set node)]
@@ -140,9 +136,9 @@
   out)
 
 (define (batch-replace b f)
-  (define out (make-mutable-batch))
+  (define out (make-batch))
   (define mapping (make-vector (batch-length b) -1))
-  (for ([node (in-vector (batch-nodes b))]
+  (for ([node (in-batch b)]
         [idx (in-naturals)])
     (define replacement (f (expr-recurse node (lambda (x) (batchref b x)))))
     (define final-idx
@@ -154,31 +150,31 @@
            (when (= -1 (vector-ref mapping idx))
              (error 'batch-replace "Replacement ~a references unknown index ~a" replacement idx))
            (vector-ref mapping idx)]
-          [_ (mutable-batch-push! out (expr-recurse expr loop))])))
+          [_ (batch-push! out (expr-recurse expr loop))])))
     (vector-set! mapping idx final-idx))
   (define roots (vector-map (curry vector-ref mapping) (batch-roots b)))
-  (mutable-batch->batch out roots))
+  (set-batch-roots! out roots)
+  out)
 
 ;; Function returns indices of alive nodes within a batch for given roots,
 ;;   where alive node is a child of a root + meets a condition - (condition node)
 (define (batch-alive-nodes batch
                            [roots (batch-roots batch)]
                            #:keep-vars-alive [keep-vars-alive #f]
                            #:condition [condition (const #t)])
-  (define nodes (batch-nodes batch))
-  (define nodes-length (batch-length batch))
-  (define alive-mask (make-vector nodes-length #f))
+  (define len (batch-length batch))
+  (define alive-mask (make-vector len #f))
   (for ([root (in-vector roots)])
     (vector-set! alive-mask root #t))
-  (for ([i (in-range (- nodes-length 1) -1 -1)]
-        [node (in-vector nodes (- nodes-length 1) -1 -1)]
-        [alv (in-vector alive-mask (- nodes-length 1) -1 -1)]
+  (for ([i (in-range (- len 1) -1 -1)]
+        [node (in-batch batch (- len 1) -1 -1)]
+        [alv (in-vector alive-mask (- len 1) -1 -1)]
         #:when (or (and alv (condition node)) (and keep-vars-alive (symbol? node))))
     (unless alv ; if keep-vars-alive then alv may not be #t, making sure it's #t
       (vector-set! alive-mask i #t))
     (expr-recurse node
                   (λ (n)
-                    (when (condition (vector-ref nodes n))
+                    (when (condition (batch-ref batch n))
                       (vector-set! alive-mask n #t)))))
   ; Return indices of alive nodes in ascending order
   (for/vector ([alv (in-vector alive-mask)]
@@ -189,7 +185,7 @@
 ;; Function constructs a vector of expressions for the given nodes of a batch
 (define (batch-reconstruct-exprs batch)
   (define exprs (make-vector (batch-length batch)))
-  (for ([node (in-vector (batch-nodes batch))]
+  (for ([node (in-batch batch)]
         [idx (in-naturals)])
     (vector-set! exprs idx (expr-recurse node (lambda (x) (vector-ref exprs x)))))
   exprs)
@@ -199,36 +195,34 @@
 ;; Space complexity: O(|N| + |N*| + |R|), where |N*| is a length of nodes without zombie nodes
 ;; The flag keep-vars is used in compiler.rkt when vars should be preserved no matter what
 (define (batch-remove-zombie batch [roots (batch-roots batch)] #:keep-vars [keep-vars #f])
-  (define nodes (batch-nodes batch))
-  (define nodes-length (batch-length batch))
-  (match (zero? nodes-length)
+  (define len (batch-length batch))
+  (match (zero? len)
     [#f
      (define alive-nodes (batch-alive-nodes batch roots #:keep-vars-alive keep-vars))
 
-     (define mappings (make-vector nodes-length -1))
+     (define mappings (make-vector len -1))
      (define (remap idx)
        (vector-ref mappings idx))
 
-     (define out (make-mutable-batch))
+     (define out (make-batch))
      (for ([alv (in-vector alive-nodes)])
-       (define node (vector-ref nodes alv))
-       (vector-set! mappings alv (mutable-batch-push! out (expr-recurse node remap))))
+       (define node (batch-ref batch alv))
+       (vector-set! mappings alv (batch-push! out (expr-recurse node remap))))
 
      (define roots* (vector-map (curry vector-ref mappings) roots))
-     (mutable-batch->batch out roots*)]
+     (set-batch-roots! out roots*)
+     out]
     [#t (batch-copy batch)]))
 
 (define (batch-ref batch reg)
+  (dvector-ref (batch-nodes batch) reg))
+
+(define (batch-pull batch reg)
   (define (unmunge reg)
-    (define node (vector-ref (batch-nodes batch) reg))
+    (define node (batch-ref batch reg))
     (expr-recurse node unmunge))
   (unmunge reg))
 
-(define (batch-restore-index batch)
-  (make-hash (for/list ([node (in-vector (batch-nodes batch))]
-                        [n (in-naturals)])
-               (cons node n))))
-
 ; Tests for progs->batch and batch->progs
 (module+ test
   (require rackunit)
@@ -253,26 +247,26 @@
 (module+ test
   (require rackunit)
   (define (zombie-test #:nodes nodes #:roots roots)
-    (define in-batch (batch nodes roots))
+    (define in-batch (batch nodes (make-hash) (make-hasheq) roots))
     (define out-batch (batch-remove-zombie in-batch))
     (check-equal? (batch->progs out-batch) (batch->progs in-batch))
     (batch-nodes out-batch))
 
-  (check-equal? (vector 0 '(sqrt 0) 2 '(pow 2 1))
-                (zombie-test #:nodes (vector 0 1 '(sqrt 0) 2 '(pow 3 2)) #:roots (vector 4)))
-  (check-equal? (vector 0 '(sqrt 0) '(exp 1))
-                (zombie-test #:nodes (vector 0 6 '(pow 0 1) '(* 2 0) '(sqrt 0) '(exp 4))
+  (check-equal? (create-dvector 0 '(sqrt 0) 2 '(pow 2 1))
+                (zombie-test #:nodes (create-dvector 0 1 '(sqrt 0) 2 '(pow 3 2)) #:roots (vector 4)))
+  (check-equal? (create-dvector 0 '(sqrt 0) '(exp 1))
+                (zombie-test #:nodes (create-dvector 0 6 '(pow 0 1) '(* 2 0) '(sqrt 0) '(exp 4))
                              #:roots (vector 5)))
-  (check-equal? (vector 0 1/2 '(+ 0 1))
-                (zombie-test #:nodes (vector 0 1/2 '(+ 0 1) '(* 2 0)) #:roots (vector 2)))
-  (check-equal? (vector 0 1/2 '(exp 1) (approx 2 0))
-                (zombie-test #:nodes (vector 0 1/2 '(+ 0 1) '(* 2 0) '(exp 1) (approx 4 0))
+  (check-equal? (create-dvector 0 1/2 '(+ 0 1))
+                (zombie-test #:nodes (create-dvector 0 1/2 '(+ 0 1) '(* 2 0)) #:roots (vector 2)))
+  (check-equal? (create-dvector 0 1/2 '(exp 1) (approx 2 0))
+                (zombie-test #:nodes (create-dvector 0 1/2 '(+ 0 1) '(* 2 0) '(exp 1) (approx 4 0))
                              #:roots (vector 5)))
-  (check-equal? (vector 'x 2 1/2 '(* 0 0) (approx 3 1) '(pow 2 4))
-                (zombie-test #:nodes
-                             (vector 'x 2 1/2 '(sqrt 1) '(cbrt 1) '(* 0 0) (approx 5 1) '(pow 2 6))
-                             #:roots (vector 7)))
-  (check-equal? (vector 'x 2 1/2 '(sqrt 1) '(* 0 0) (approx 4 1) '(pow 2 5))
-                (zombie-test #:nodes
-                             (vector 'x 2 1/2 '(sqrt 1) '(cbrt 1) '(* 0 0) (approx 5 1) '(pow 2 6))
-                             #:roots (vector 7 3))))
+  (check-equal?
+   (create-dvector 'x 2 1/2 '(* 0 0) (approx 3 1) '(pow 2 4))
+   (zombie-test #:nodes (create-dvector 'x 2 1/2 '(sqrt 1) '(cbrt 1) '(* 0 0) (approx 5 1) '(pow 2 6))
+                #:roots (vector 7)))
+  (check-equal?
+   (create-dvector 'x 2 1/2 '(sqrt 1) '(* 0 0) (approx 4 1) '(pow 2 5))
+   (zombie-test #:nodes (create-dvector 'x 2 1/2 '(sqrt 1) '(cbrt 1) '(* 0 0) (approx 5 1) '(pow 2 6))
+                #:roots (vector 7 3))))

src/core/compiler.rkt

@@ -73,7 +73,7 @@
 
   (define instructions
     (for/vector #:length (- (batch-length batch*) num-vars)
-                ([node (in-vector (batch-nodes batch*) num-vars)])
+                ([node (in-batch batch* num-vars)])
       (match node
         [(literal value (app get-representation repr)) (list (const (real->repr value repr)))]
         [(list op args ...) (cons (impl-info op 'fl) args)])))