Merge pull request #771 from AayushSabharwal/as/poly-variant

feat: add `Polyform` variant, remove `AddOrMul`, `Pow`

Author: AayushSabharwal

Project.toml

@@ -9,6 +9,7 @@ ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 Bijections = "e2ed5e7c-b2de-5872-ae92-c73ca462fb04"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
+ConcurrentUtilities = "f0e56b4a-5159-44fe-b623-3e5288b988bb"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
@@ -20,6 +21,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Moshi = "2e0e35c7-a2e4-4343-998d-7ef72827ed2d"
 MultivariatePolynomials = "102ac46a-7ee4-5c85-9060-abc95bfdeaa3"
+MutableArithmetics = "d8a4904e-b15c-11e9-3269-09a3773c0cb0"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 ReadOnlyArrays = "988b38a3-91fc-5605-94a2-ee2116b3bd83"
 ReadOnlyDicts = "795d4caa-f5a7-4580-b5d8-c01d53451803"
@@ -48,8 +50,9 @@ ArrayInterface = "7.8"
 Bijections = "0.1.2, 0.2"
 ChainRulesCore = "1"
 Combinatorics = "1.0"
+ConcurrentUtilities = "2.5.0"
 ConstructionBase = "1.5.7"
-DataStructures = "0.18"
+DataStructures = "0.19"
 DocStringExtensions = "0.8, 0.9"
 DynamicPolynomials = "0.5, 0.6"
 EnumX = "1.0.5"
@@ -58,6 +61,7 @@ LabelledArrays = "1.5"
 MacroTools = "0.5.16"
 Moshi = "0.3.6"
 MultivariatePolynomials = "0.5"
+MutableArithmetics = "1.6.4"
 NaNMath = "0.3, 1.1.2"
 OhMyThreads = "0.7"
 ReadOnlyArrays = "0.2.0"
@@ -91,3 +95,8 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
 test = ["BenchmarkTools", "Documenter", "LabelledArrays", "Pkg", "PkgBenchmark", "Random", "ReferenceTests", "ReverseDiff", "SafeTestsets", "Test", "Zygote", "OhMyThreads", "RuntimeGeneratedFunctions"]
+
+[sources]
+DynamicPolynomials = {url = "https://github.com/AayushSabharwal/DynamicPolynomials.jl", rev = "as/new-poly-merge"}
+MultivariatePolynomials = {url = "https://github.com/AayushSabharwal/MultivariatePolynomials.jl", rev = "as/poly-merge-nonconcrete"}
+MutableArithmetics = {url = "https://github.com/AayushSabharwal/MutableArithmetics.jl", rev="as+bl/simplify_promote_type_fallback"}

src/SymbolicUtils.jl

@@ -31,9 +31,71 @@ using WeakValueDicts: WeakValueDict
 using WeakCacheSets: WeakCacheSet, getkey!
 using Base: RefValue
 import MacroTools
+import MultivariatePolynomials as MP
+import DynamicPolynomials as DP
+import MutableArithmetics as MA
+import ConcurrentUtilities: ReadWriteLock, readlock, readunlock
 
 # include("WeakCacheSets.jl")
 
+macro manually_scope(val, expr, is_forced = false)
+    @assert Meta.isexpr(val, :call)
+    @assert val.args[1] == :(=>)
+
+    var_name = val.args[2]
+    new_val = val.args[3]
+    old_name = gensym(:old_val)
+    cur_name = gensym(:cur_val)
+    retval_name = gensym(:retval)
+    close_expr = :($var_name[] = $old_name)
+    interpolated_expr = MacroTools.postwalk(expr) do ex
+        if Meta.isexpr(ex, :$) && length(ex.args) == 1 && ex.args[1] == :$
+            return cur_name
+        else
+            return ex
+        end
+    end
+    basic_result = quote
+        $cur_name = $var_name[] = $new_val
+        $retval_name = try
+            $interpolated_expr
+        finally
+            $close_expr
+        end
+    end
+    is_forced && return quote
+        $old_name = $var_name[]
+        $basic_result
+    end |> esc
+
+    return quote
+        $old_name = $var_name[]
+        if $iszero($old_name)
+            $basic_result
+        else
+            $cur_name = $old_name
+            $retval_name = begin
+                $interpolated_expr
+            end
+        end
+        $retval_name
+    end |> esc
+end
+
+# copied from https://github.com/JuliaLang/julia/blob/80f7db8e51b2ba1dd21e913611c23a6d5b75ecab/base/lock.jl#L371-L381
+# and adapted for readlock/readunlock
+macro readlock(l, expr)
+    quote
+        temp = $(esc(l))
+        $readlock(temp)
+        try
+            $(esc(expr))
+        finally
+            $readunlock(temp)
+        end
+    end
+end
+
 include("cache.jl")
 Base.@deprecate istree iscall
 
@@ -66,7 +128,6 @@ include("matchers.jl")
 include("rewriters.jl")
 
 # Convert to an efficient multi-variate polynomial representation
-import MultivariatePolynomials as MP
 import DynamicPolynomials
 export expand
 include("polyform.jl")