git » blacktide.git » commit 962feec

diff --git a/src/com/tailrecursion/blacktide/compiler/passes/source.clj b/src/com/tailrecursion/blacktide/compiler/passes/source.clj
index 448bf47..874510a 100644
--- a/src/com/tailrecursion/blacktide/compiler/passes/source.clj
+++ b/src/com/tailrecursion/blacktide/compiler/passes/source.clj
@@ -232,7 +232,8 @@
           (let [[supply inner bindings] (parse-bindings supply ctx :immutable (second selector))
                 [supply selector] (parse-operand supply inner (first (nnext selector)))
                 [supply arms] (parse-arms supply inner arms)
-                body (expression ctx {:hint nil :var-meta {}} {:op :case :selector selector :arms arms})]
+                body (expression ctx {:hint nil :var-meta {}}
+                                 {:op :case :selector (cases/compose-selector (count arms) selector) :arms arms})]
             [supply (expression ctx {:hint nil :var-meta {}} {:op :let :bindings bindings :body body})])))))
 
   (parse-dispatch
diff --git a/src/com/tailrecursion/blacktide/compiler/passes/source_case.clj b/src/com/tailrecursion/blacktide/compiler/passes/source_case.clj
index 2427814..399e1fd 100644
--- a/src/com/tailrecursion/blacktide/compiler/passes/source_case.clj
+++ b/src/com/tailrecursion/blacktide/compiler/passes/source_case.clj
@@ -5,6 +5,25 @@
             [com.tailrecursion.blacktide.compiler.languages :as l]
             [typed.clojure :as-alias t]))
 
+(np/defpass compose-selector
+  {:from [l/Bound Expr] :to [l/Bound Expr] :args [t/Int] :body true}
+  [arm-count expr]
+  ;; parse-case consumes this result only as a source-arm ordinal. Compose the
+  ;; mappings without changing constants, collision checks or source bodies.
+  (let [node (:node expr)]
+    (if (and (= :case (:op node)) (= :dispatch (:op (:node (:selector node)))))
+      (if-let [ordinals
+               (reduce (fn [^{::t/- (t/Option (t/Vec t/Int))} out ^{::t/- l/Bound-Expr} arm]
+                         (let [node (:node arm) value (:value node)]
+                           (if (and out (= :literal (:op node)) (integer? value) (<= 0 value) (< value arm-count))
+                             (conj out (long value)) (reduced nil)))) [] (:arms node))]
+        (assoc expr :node
+          (update (:node (:selector node)) :spec
+            #(-> % (update :default (fn [index] (nth ordinals index)))
+                   (update :tests update-vals (fn [[test index]] [test (nth ordinals index)])))))
+        expr)
+      expr)))
+
 (np/defpass expand-case
   {:from [l/Bound Forms] :to [l/Bound CaseSyntax] :args [(t/Map t/Sym t/Any)] :body true}
   [env forms]
diff --git a/test/com/tailrecursion/blacktide/case_test.clj b/test/com/tailrecursion/blacktide/case_test.clj
index 90767f7..bc23afc 100644
--- a/test/com/tailrecursion/blacktide/case_test.clj
+++ b/test/com/tailrecursion/blacktide/case_test.clj
@@ -7,6 +7,7 @@
             [clojure.test.check.properties :as prop]
             [clojure.walk :as walk]
             [com.tailrecursion.blacktide :refer [loophole setq psetq label goto]]
+            [com.tailrecursion.blacktide.compiler.passes.lower :as lower]
             [com.tailrecursion.blacktide.jvm-fixture :as jvm]
             [com.tailrecursion.blacktide.backend-fixture :as backend]))
 
@@ -130,6 +131,19 @@
   (is (= :invalid-recur (:kind (backend/rejection '(loophole [] (loop [] (inc (case 0 0 (recur) 1)))))))))
 
 (deftest prepared-dispatch-and-expansion-equivalence
+  (let [form '(loophole [x 0 out []]
+                (label again)
+                (case (do (swap! events conj x) x)
+                  2 (goto done)
+                  (0 1) (setq out (conj out x) x (inc x))
+                  (throw (Exception. "unexpected default")))
+                (goto again)
+                (label done) out)
+        graph (backend/capture #'lower/stage form)]
+    (is (= 1 (count (filter #(= :switch (:op (:term %))) (vals (:blocks graph)))))
+        "Compiler-generated ordinal selection goes directly to the source arms")
+    (doseq [mode [:default :dispatcher]]
+      (is (= {:value [0 1] :trace [0 1 2]} (backend/execution form mode)))))
   (doseq [body '[(case input 0 :zero 2 :two :other)
                 (case input "Aa" :aa "BB" :bb :other)
                 (case input (:a :b) :group nil :nil :other)