saulshanabrook
commented
5 months ago @ztatlock and others had wondered whether this PR could be simplified by not supporting curried functions. This would remove the storage of values in the sort, making the function sort not a container sort.
My sense from implementing it below is that it would be possible using codegen/macros, but it expands the amount of code and would require adding implementations for every combination of type of curried arguments for every function.
As currently implemented, this PR would allow you to define a map function over a custom user defined map data structure:
(datatype Math
(Num i64)
(Var String)
(Add Math Math)
(Mul Math Math))
(rewrite (Mul (Num x) (Num y)) (Num (* x y)))
(datatype MathList
(Nil)
(Cons Math MathList))
(sort MathFn (Fn (Math) Math))
(function list-map-math (MathFn MathList) MathList)
(rewrite (list-map-math fn (Nil)) (Nil))
(rewrite (list-map-math fn (Cons x xs)) (Cons (app fn x) (list-map-math fn xs)))
;; version where partial application is allowed
(function list-multiply-by (MathList Math) MathList)
(rewrite (list-multiply-by l i) (list-map-math (fn "Mul" i) l))However, if partial application was not allowed, we could implement this by defining a custom map function that would partially apply an arg that is passed into it in the function invocation:
;; alternative version where no partial application is allowed
;; Must make a new map function that also takes a binary fn instead of a unary one
;; as well as an initial partialled applied arg
(sort MathMathFn (Fn (Math Math) Math))
(function list-map-math-partial (MathMathFn MathList Math) MathList)
(rewrite (list-map-math-partial fn (Nil)) (Nil))
(rewrite (list-map-math-partial fn (Cons x xs) y) (Cons (app fn y x) (list-map-math-partial fn xs y)))
(rewrite (list-multiply-by l i) (list-map-math-partial (fn "Mul") l i))Or alternatively, we could write a wrapper for a unary math function that supports partial application and use this (thanks @yihozhang for this version):
(sort MathMathFn (Fn (Math Math) Math))
(datatype MathFnWrapper (partial-app MathMathFn Math))
(function App (MathFnWrapper Math) Math)
(rewrite (App (partial-app f a) b) (app f a b))
;; rewrite
(function list-map-math (MathFnWrapper MathList) MathList)
(rewrite (list-map-math fn (Nil)) (Nil))
(rewrite (list-map-math fn (Cons x xs)) (Cons (App fn x) (list-map-math fn xs)))
(rewrite (list-multiply-by l i) (list-map-math (partial-app (fn "Mul") i)))
oflatt
This PR is an attempt to support higher-order functions in egglog.
Compared to the earlier approach, in #155, this doesn't embed the bodies of unevaluated functions in the EGraph. Instead, all function values are just represented as lookups to egglog functions. An advantage of this approach is that we don't need to re-implement function application in anyway, we can just re-use our existing implementation. For defining the body of functions, we can just use rewrite rules.
This was inspired by @yihozhang's comment a while ago about how to embed functions in egglog.
As implemented, this supports functions that return primitives and eqsorts, since it uses the builtin mechanism for calling functions.
Although it doesn't support anonymous functions at the moment, adding them would be relatively trivial I believe as syntactic sugar, that would translate to a new global function with a random name and a rewrite rule to cover all cases.
I made sure that this implementation supports partial application, which is necessary for rewrite rules that transform functions.
The implementation is fairly self-contained as a primitive sort, requiring no changes to the parser or type checker. The only main breaking change is requiring a mutable EGraph pointer to be passed to primitive functions. This is necessary so that the function evaluation can possibly mutate the unionfind data structure, if new e-classes are created. I am unable to get access to a mutable egraph in the query checker, so these will panic if attempted to use in facts, they can only be used in action. This limitation seems reasonable.
One of the driving use cases for this on my end, besides just generally being cool to be able to optimize functional programs, is @sklam's example for defining a higher level array API and translating that within egglog.