Merge pull request #180 from jolt-lang/spike/fast-arith-return-hints

Numeric return-type hints (round 3)

host/chez/host-contract.ss

@@ -32,6 +32,9 @@
 (define hc-kw-name  (keyword #f "name"))
 (define hc-kw-var   (keyword #f "var"))
 (define hc-kw-unresolved (keyword #f "unresolved"))
+(define hc-kw-num-ret (keyword #f "num-ret"))
+(define hc-kw-double (keyword #f "double"))
+(define hc-kw-long (keyword #f "long"))
 (define hc-kw-regex (keyword #f "regex"))
 (define hc-kw-inst  (keyword #f "#inst"))
 (define hc-kw-uuid  (keyword #f "#uuid"))
@@ -176,13 +179,24 @@
 ;; No :host branch: there is no separate native-op env — the hot
 ;; clojure.core primitives (+,-,map,...) are declared in clojure.core below so
 ;; they classify as :var and the emitter's native-op path lowers them.
+;; A var's declared numeric return (^double/^long on its name) -> :double/:long,
+;; read from its meta. Lets jolt.passes.numeric type a call to it.
+(define (hc-cell-num-ret cell)
+  (let ((m (and cell (hashtable-ref var-meta-table cell #f))))
+    (and m (let ((t (jolt-get m hc-kw-tag)))
+             (cond ((equal? t "double") hc-kw-double)
+                   ((equal? t "long") hc-kw-long)
+                   (else #f))))))
+
 (define (hc-resolve-global ctx sym)
   (let* ((nm (symbol-t-name sym))
          (cell (hc-resolve-cell ctx sym)))
     (if (and cell (var-cell-defined? cell))
-        (jolt-hash-map hc-kw-kind hc-kw-var
-                       hc-kw-ns (var-cell-ns cell)
-                       hc-kw-name (var-cell-name cell))
+        (let ((base (jolt-hash-map hc-kw-kind hc-kw-var
+                                   hc-kw-ns (var-cell-ns cell)
+                                   hc-kw-name (var-cell-name cell)))
+              (nr (hc-cell-num-ret cell)))
+          (if nr (jolt-assoc base hc-kw-num-ret nr) base))
         (jolt-hash-map hc-kw-kind hc-kw-unresolved hc-kw-name nm))))
 
 (define (hc-intern! ctx ns-name nm) (declare-var! ns-name nm) jolt-nil)

jolt-core/jolt/analyzer.clj

@@ -80,15 +80,26 @@
   (let [m (form-sym-meta sym)]
     (when m (let [t (get m :tag)] (when t (record-ctor-key ctx t))))))
 
-;; A primitive numeric hint (^long / ^double) on a binding symbol -> :long /
-;; :double, else nil. Drives the fl*/fx* fast path (jolt.passes.numeric). The tag
-;; is a string on the data reader; tolerate a symbol from macroexpansion too.
-(defn- nhint-of [ctx sym]
-  (let [m (form-sym-meta sym)
-        t (when m (get m :tag))
-        s (cond (form-sym? t) (form-sym-name t) (string? t) t :else nil)]
+;; A ^long / ^double tag -> :long / :double, else nil. The tag is a string on the
+;; data reader; tolerate a symbol from macroexpansion too.
+(defn- tag->nkind [t]
+  (let [s (cond (form-sym? t) (form-sym-name t) (string? t) t :else nil)]
     (cond (= s "double") :double (= s "long") :long :else nil)))
 
+;; A primitive numeric hint (^long / ^double) on a binding symbol. Drives the
+;; fl*/fx* fast path (jolt.passes.numeric).
+(defn- nhint-of [ctx sym]
+  (let [m (form-sym-meta sym)] (when m (tag->nkind (get m :tag)))))
+
+;; Push a numeric return hint (from ^double/^long on a defn's name) onto each arity
+;; of its fn, so the back end coerces the body's value to that kind on return —
+;; making the hint a contract a caller's arithmetic can trust.
+(defn- with-ret-nhint [node kind]
+  (if (and kind (= :fn (:op node)))
+    (assoc node :arities (mapv (fn [a] (if (:ret-nhint a) a (assoc a :ret-nhint kind)))
+                               (:arities node)))
+    node))
+
 (defn- analyze-seq [ctx forms env]
   (let [v (mapv #(analyze ctx % env) forms)
         n (count v)]
@@ -304,7 +315,8 @@
                         base0)
             node-meta (if has-doc (assoc base-meta :doc (nth items 2)) base-meta)]
         (host-intern! ctx cur nm)
-        (def-node cur nm (analyze ctx val-form env) node-meta)))))
+        ;; a ^double/^long return hint on the name applies to all arities of the fn.
+        (def-node cur nm (with-ret-nhint (analyze ctx val-form env) (tag->nkind tag)) node-meta)))))
 
 ;; (set! (.-field obj) v) mutates a deftype instance field in place; (set! *var* v)
 ;; sets the var's innermost thread binding, else its root. A local target (jolt
@@ -493,14 +505,16 @@
         (let [h (get (:hints env) nm)] (if h (assoc (local nm) :hint h) (local nm)))
       ns (let [r (resolve-global ctx form)]
            (if (= :var (:kind r))
-             (var-ref (:ns r) (:name r))
+             (cond-> (var-ref (:ns r) (:name r)) (:num-ret r) (assoc :num-ret (:num-ret r)))
              ;; A non-var qualified ref `Class/member` is a host class static
              ;; (Math/sqrt, Long/MAX_VALUE, System/getenv). The Chez back end
              ;; lowers it to a runtime static dispatch.
              (host-static ns nm)))
       :else (let [r (resolve-global ctx form)]
               (case (:kind r)
-                :var (var-ref (:ns r) (:name r))
+                ;; :num-ret (a ^double/^long declared return) rides on the var node so
+                ;; jolt.passes.numeric types a call to it (an accumulator over the result).
+                :var (cond-> (var-ref (:ns r) (:name r)) (:num-ret r) (assoc :num-ret (:num-ret r)))
                 :host (host-ref (:name r))
                 ;; :unresolved — emitting a var-ref here would auto-intern an
                 ;; UNBOUND var, so a typo'd symbol would die later as 'Cannot call

jolt-core/jolt/backend_scheme.clj

@@ -373,8 +373,15 @@
         pbind (map (fn [o p] (str "(" p " " (nhint-init nh o p) ")")) orig params)
         binds (if restp
                 (concat pbind [(str "(" restp " (list->cseq " restp "))")])
-                pbind)]
-    [paramlist (str "(let " label " (" (str/join " " binds) ") " body ")")]))
+                pbind)
+        lett (str "(let " label " (" (str/join " " binds) ") " body ")")
+        ;; a ^double/^long return hint coerces the arity's value on the way out
+        ;; (exact->inexact / jolt->fx), like a JVM primitive return — so a caller's
+        ;; arithmetic over the result is sound.
+        ret (:ret-nhint a)]
+    [paramlist (cond (= ret :double) (str "(exact->inexact " lett ")")
+                     (= ret :long)   (str "(jolt->fx " lett ")")
+                     :else lett)]))
 
 (defn- emit-fn [node]
   (let [arities (:arities node)

jolt-core/jolt/passes/numeric.clj

@@ -100,8 +100,13 @@
         argnodes (mapv (fn [r] (nth r 1)) ars)
         node1 (assoc node :args argnodes)
         n (count ars)]
-    (if (nil? nm)
-      [nil node1]
+    (cond
+      ;; a call to a var with a declared numeric return (^double/^long) yields that
+      ;; kind, so an accumulator over the result types. The call itself isn't an
+      ;; arithmetic op to lower — its body already coerces the return.
+      (get fnode :num-ret) [(get fnode :num-ret) node1]
+      (nil? nm) [nil node1]
+      :else
       (let [;; per-operand class: :double / :long (typed), :wild (integer literal,
             ;; usable in either), or :no (anything else — blocks specialization).
             cls (mapv (fn [r] (let [k (nth r 0) nd (nth r 1)]

test/chez/numeric-test.ss

@@ -100,5 +100,24 @@
 (ok "long-seeded loop accumulator counts to 100"
     (= 100 (jnum->exact (ev "((fn* ([^long start] (loop [acc start] (if (< acc 100) (recur (inc acc)) acc)))) 0)"))))
 
+;; --- numeric return hints (round 3) ---
+;; a ^double / ^long return hint coerces the fn's value on the way out (the contract).
+(jolt-compile-eval "(def ^double dsq (fn* ([x] (* x x))))" "u")
+(ok "^double return coerces value to a flonum" (flonum? (ev "(dsq 3)")))
+(jolt-compile-eval "(def ^long ldbl (fn* ([x] (+ x x))))" "u")
+(ok "^long return coerces value to a fixnum" (let ((r (ev "(ldbl 5)"))) (and (fixnum? r) (= r 10))))
+;; the defn macro must carry the name's return hint through to the def.
+(jolt-compile-eval "(defn ^double dnsq [x] (* x x))" "u")
+(ok "defn ^double return coerces to flonum" (flonum? (ev "(dnsq 4)")))
+
+;; caller propagation: a call to a ^double-returning fn types an accumulator over it.
+(let ((e (emitf "u" "(fn* ([] (loop [acc 0.0 i 0] (if (< i 3) (recur (+ acc (dsq 2.0)) (inc i)) acc))))")))
+  (ok "accumulator over a ^double-returning call lowers to fl+" (has? e "(fl+")))
+(ok "accumulator over ^double call runtime: 3 * (2*2) = 12.0"
+    (= 12 (jnum->exact (ev "((fn* ([] (loop [acc 0.0 i 0] (if (< i 3) (recur (+ acc (dsq 2.0)) (inc i)) acc)))))"))))
+;; a ^double call result also specializes a straight-line op
+(let ((e (emitf "u" "(fn* ([^double y] (+ y (dsq 2.0))))")))
+  (ok "straight-line op over ^double call lowers to fl+" (has? e "(fl+")))
+
 (printf "~a/~a passed~n" (- total fails) total)
 (exit (if (zero? fails) 0 1))