riddley.walk.clj Maven / Gradle / Ivy
(ns riddley.walk
(:refer-clojure :exclude [macroexpand])
(:require
[riddley.compiler :as cmp]))
(defn macroexpand
"Expands both macros and inline functions. Optionally takes a `special-form?` predicate which
identifies first elements of expressions that shouldn't be macroexpanded, and honors local
bindings."
([x]
(macroexpand x nil))
([x special-form?]
(cmp/with-base-env
(if (seq? x)
(let [frst (first x)]
(if (or
(and special-form? (special-form? frst))
(contains? (cmp/locals) frst))
;; might look like a macro, but for our purposes it isn't
x
(let [x' (macroexpand-1 x)]
(if-not (identical? x x')
(macroexpand x' special-form?)
;; if we can't macroexpand any further, check if it's an inlined function
(if-let [inline-fn (and (seq? x')
(symbol? (first x'))
(-> x' meta ::transformed not)
(or
(-> x' first resolve meta :inline-arities not)
((-> x' first resolve meta :inline-arities)
(count (rest x'))))
(-> x' first resolve meta :inline))]
(let [x'' (apply inline-fn (rest x'))]
(macroexpand
;; unfortunately, static function calls can look a lot like what we just
;; expanded, so prevent infinite expansion
(if (= '. (first x''))
(with-meta
(concat (butlast x'')
[(if (instance? clojure.lang.IObj (last x''))
(with-meta (last x'')
(merge
(meta (last x''))
{::transformed true}))
(last x''))])
(meta x''))
x'')
special-form?))
x')))))
x))))
;;;
(defn- do-handler [f [_ & body]]
(list* 'do
(doall
(map f body))))
(defn- fn-handler [f x]
(let [prelude (take-while (complement sequential?) x)
remainder (drop (count prelude) x)
remainder (if (vector? (first remainder))
(list remainder) remainder)
body-handler (fn [x]
(cmp/with-lexical-scoping
(doseq [arg (first x)]
(cmp/register-arg arg))
(doall
(list* (first x)
(map f (rest x))))))]
(cmp/with-lexical-scoping
;; register a local for the function, if it's named
(when-let [nm (second prelude)]
(cmp/register-local nm
(list* 'fn* nm
(map #(take 1 %) remainder))))
(concat
prelude
(if (seq? (first remainder))
(doall (map body-handler remainder))
[(body-handler remainder)])))))
(defn- def-handler [f x]
(let [[_ n & r] x]
(cmp/with-lexical-scoping
(cmp/register-local n '())
(list* 'def (f n) (doall (map f r))))))
(defn- let-bindings [f x]
(->> x
(partition-all 2)
(mapcat
(fn [[k v]]
(let [[k v] [k (f v)]]
(cmp/register-local k v)
[k v])))
vec))
(defn- reify-handler [f x]
(let [[_ classes & fns] x]
(list* 'reify* classes
(doall
(map
(fn [[nm args & body]]
(cmp/with-lexical-scoping
(doseq [arg args]
(cmp/register-arg arg))
(list* nm args (doall (map f body)))))
fns)))))
(defn- deftype-handler [f x]
(let [[_ type resolved-type args _ interfaces & fns] x]
(cmp/with-lexical-scoping
(doseq [arg args]
(cmp/register-arg arg))
(list* 'deftype* type resolved-type args :implements interfaces
(doall
(map
(fn [[nm args & body]]
(cmp/with-lexical-scoping
(doseq [arg args]
(cmp/register-arg arg))
(list* nm args (doall (map f body)))))
fns))))))
(defn- let-handler [f x]
(cmp/with-lexical-scoping
(doall
(list*
(first x)
(let-bindings f (second x))
(map f (drop 2 x))))))
(defn- case-handler [f x]
(let [prefix (butlast (take-while (complement map?) x))
default (last (take-while (complement map?) x))
body (first (drop-while (complement map?) x))
suffix (rest (drop-while (complement map?) x))]
(concat
prefix
[(f default)]
[(let [m (->> body
(map
(fn [[k [idx form]]]
[k [idx (f form)]]))
(into {}))]
(if (every? number? (keys m))
(into (sorted-map) m)
m))]
suffix)))
(defn- catch-handler [f x]
(let [[_ type var & body] x]
(cmp/with-lexical-scoping
(when var
(cmp/register-arg (with-meta var (merge (meta var) {:tag type}))))
(list* 'catch type var
(doall (map f body))))))
(defn- dot-handler [f x]
(let [[_ hostexpr mem-or-meth & remainder] x]
(list* '.
(f hostexpr)
(if (seq? mem-or-meth)
(list* (first mem-or-meth)
(doall (map f (rest mem-or-meth))))
(f mem-or-meth))
(doall (map f remainder)))))
(defn walk-exprs
"A walk function which only traverses valid Clojure expressions. The `predicate` describes
whether the sub-form should be transformed. If it returns true, `handler` is invoked, and
returns a transformed form.
Unlike `clojure.walk`, if the handler is called, the rest of the sub-form is not walked.
The handler function is responsible for recursively calling `walk-exprs` on the form it is
given.
Macroexpansion can be halted by defining a set of `special-form?` which will be left alone.
Including `fn`, `let`, or other binding forms can break local variable analysis, so use
with caution."
([predicate handler x]
(walk-exprs predicate handler nil x))
([predicate handler special-form? x]
(cmp/with-base-env
(let [x (try
(macroexpand x special-form?)
(catch ClassNotFoundException _
x))
walk-exprs' (partial walk-exprs predicate handler special-form?)
x' (cond
(and (seq? x) (= 'var (first x)) (predicate x))
(handler (eval x))
(and (seq? x) (= 'quote (first x)) (not (predicate x)))
x
(predicate x)
(handler x)
(seq? x)
((condp = (first x)
'do do-handler
'def def-handler
'fn* fn-handler
'let* let-handler
'loop* let-handler
'letfn* let-handler
'case* case-handler
'catch catch-handler
'reify* reify-handler
'deftype* deftype-handler
'. dot-handler
#(doall (map %1 %2)))
walk-exprs' x)
(instance? java.util.Map$Entry x)
(clojure.lang.MapEntry.
(walk-exprs' (key x))
(walk-exprs' (val x)))
(or
(set? x)
(vector? x))
(into (empty x) (map walk-exprs' x))
(instance? clojure.lang.IRecord x)
x
(map? x)
(into (empty x) (map walk-exprs' x))
;; special case to handle clojure.test
(and (symbol? x) (-> x meta :test))
(vary-meta x update-in [:test] walk-exprs')
:else
x)]
(if (instance? clojure.lang.IObj x')
(with-meta x' (merge (meta x) (meta x')))
x')))))
;;;
(defn macroexpand-all
"Recursively macroexpands all forms, preserving the &env special variables."
[x]
(walk-exprs (constantly false) nil x))