git » blacktide.git » master » tree

[master] / src / com / tailrecursion / blacktide / compiler / passes / source_case.clj

(ns ^:typed.clojure com.tailrecursion.blacktide.compiler.passes.source-case
  "Ask core case to select source-arm ordinals. Constants stay opaque, while
  source bodies remain unique. Decode prepared case* tables without evaluating them."
  (:require [com.tailrecursion.nanopass :as np]
            [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]
  (let [_ (when (< (count forms) 2)
            (.applyTo ^clojure.lang.Var #'clojure.core/case (cons (apply list forms) (cons env (rest forms)))))
        clauses (subvec forms 2) pairs (vec (partition 2 clauses)) default? (odd? (count clauses))
        arms (cond-> (mapv second pairs) default? (conj (peek clauses)))
        clauses (into [] (mapcat (fn [^{::t/- t/Int} index] [(first (nth pairs index)) index])) (range (count pairs)))
        clauses (cond-> clauses default? (conj (count pairs)))
        form (list* 'clojure.core/case (second forms) clauses)]
    {:selector (.applyTo ^clojure.lang.Var #'clojure.core/case (cons form (cons env (rest form)))) :arms arms}))

(np/defpass int32 {:from [l/Bound Raw] :to t/Int :body true}
  [value]
  (if (instance? Number value) (.intValue ^Number value)
      (throw (IllegalArgumentException. "case* expects numeric dispatch data"))))

(np/defpass dispatch-syntax {:from [l/Bound Forms] :to [l/Bound DispatchSyntax] :body true}
  [forms]
  (let [_ (when-not (<= 8 (count forms) 9) (throw (IllegalArgumentException. "Malformed case*")))
        [_ value shift mask default table switch test skip] forms]
    (when-not (and (map? table) (seq table) (#{:compact :sparse} switch)
                   (#{:int :hash-equiv :hash-identity} test) (or (nil? skip) (set? skip)))
      (throw (IllegalArgumentException. "Malformed case* dispatch table")))
    (let [[tests arms]
          (reduce (fn [^{::t/- (t/HVec [(t/Map t/Int (t/HVec [t/Any t/Int])) l/Bound-Forms])} acc
                       ^{::t/- (t/HVec [t/Any t/Any])} entry]
                    (let [[tests arms] acc [key pair] entry]
                      (when-not (and (vector? pair) (= 2 (count pair)))
                        (throw (IllegalArgumentException. "case* table entries require constant/body pairs")))
                      [(assoc tests (int32 key) [(if (= :int test) (int32 (first pair)) (first pair)) (count arms)])
                       (conj arms (second pair))])) [{} []] (mapv (fn [[key value]] [key value]) table))]
      {:value value :arms (conj arms default)
       :spec {:shift (int32 shift) :mask (int32 mask) :default (count arms) :tests tests
              :switch switch :test test :skip (when skip (set (map int32 skip)))}})))