Merge pull request #761 from JuliaSymbolics/as/const-variant

[WIP] feat: add `Const` variant

Author: AayushSabharwal

src/SymbolicUtils.jl

@@ -106,7 +106,7 @@ Base.@deprecate istree iscall
 
 include("small_array.jl")
 
-export istree, operation, arguments, sorted_arguments, iscall
+export istree, operation, arguments, sorted_arguments, iscall, unwrap_const
 # Sym, Term,
 # Add, Mul and Pow
 include("types.jl")

src/code.jl

@@ -10,7 +10,8 @@ export toexpr, Assignment, (←), Let, Func, DestructuredArgs, LiteralExpr,
 import ..SymbolicUtils
 import ..SymbolicUtils.Rewriters
 import SymbolicUtils: @matchable, BasicSymbolic, Sym, Term, iscall, operation, arguments, issym,
-                      symtype, sorted_arguments, metadata, isterm, term, maketerm, Symbolic
+                      symtype, sorted_arguments, metadata, isterm, term, maketerm, unwrap_const,
+                      ArgsT, maybe_const
 import SymbolicIndexingInterface: symbolic_type, NotSymbolic
 
 ##== state management ==##
@@ -142,17 +143,20 @@ function function_to_expr(op::Union{typeof(*),typeof(+)}, O, st)
 end
 
 function function_to_expr(op::typeof(^), O, st)
-    args = arguments(O)
-    if args[2] isa Real && args[2] < 0
-        args[1] = Term(inv, Any[args[1]])
-        args[2] = -args[2]
-    end
-    if isequal(args[2], 1)
-        return toexpr(args[1], st)
+    base, exp = arguments(O)
+    base = unwrap_const(base)
+    exp = unwrap_const(exp)
+    if exp isa Real && exp < 0
+        base = Term(inv, ArgsT((base,)))
+        if isone(-exp)
+            return toexpr(base, st)
+        else
+            exp = -exp
+        end
     end
-    if get(st.rewrites, :nanmath, false) === true && !(args[2] isa Integer)
+    if get(st.rewrites, :nanmath, false) === true && !(exp isa Integer)
         op = NaNMath.pow
-        return toexpr(Term(op, args), st)
+        return toexpr(Term(op, ArgsT((maybe_const(base), maybe_const(exp)))), st)
     end
     return nothing
 end
@@ -203,11 +207,11 @@ end
 _is_tuple_of_symbolics(O) = false
 
 function toexpr(O, st)
+    O = unwrap_const(O)
+    O = substitute_name(O, st)
     if issym(O)
-        O = substitute_name(O, st)
-        return issym(O) ? nameof(O) : toexpr(O, st)
+        return nameof(O)
     end
-    O = substitute_name(O, st)
 
     if _is_array_of_symbolics(O)
         return issparse(O) ? toexpr(MakeSparseArray(O)) : toexpr(MakeArray(O, typeof(O)), st)
@@ -908,6 +912,7 @@ function cse!(expr::BasicSymbolic, state::CSEState)
         args = arguments(expr)
         cse_inside_expr(expr, op, args...) || return expr
         args = map(args) do arg
+            arg = unwrap_const(arg)
             if arg isa Union{Tuple, AbstractArray} &&
                 (_is_array_of_symbolics(arg) || _is_tuple_of_symbolics(arg))
                 if arg isa Tuple

src/inspect.jl

@@ -1,9 +1,6 @@
 import AbstractTrees
 
 const inspect_metadata = Ref{Bool}(false)
-function AbstractTrees.nodevalue(x::Symbolic)
-    iscall(x) ? operation(x) : isexpr(x) ? head(x) : x
-end
 
 function AbstractTrees.nodevalue(x::BSImpl.Type)
     T = nameof(MData.variant_type(x))
@@ -35,7 +32,7 @@ the expression.
 
 This function is used internally for printing via AbstractTrees.
 """
-function AbstractTrees.children(x::Symbolic)
+function AbstractTrees.children(x::BasicSymbolic)
     iscall(x) ? sorted_arguments(x) : isexpr(x) ? sorted_children(x) : ()
 end
 
@@ -50,7 +47,7 @@ Line numbers will be shown, use `pluck(expr, line_number)` to get the sub expres
 """
 function inspect end
 
-function inspect(io::IO, x::Symbolic;
+function inspect(io::IO, x::BasicSymbolic;
         hint=true,
         metadata=inspect_metadata[])
 

src/matchers.jl

@@ -7,11 +7,12 @@
 #
 
 function matcher(val::Any, acSets)
+    val = unwrap_const(val)
     # if val is a call (like an operation) creates a term matcher or term matcher with defslot
     if iscall(val)
         # if has two arguments and one of them is a DefSlot, create a term matcher with defslot
         # just two arguments bc defslot is only supported with operations with two args: *, ^, +
-        if any(x -> isa(x, DefSlot), parent(arguments(val)))
+        if any(x -> isa(unwrap_const(x), DefSlot), parent(arguments(val)))
             return defslot_term_matcher_constructor(val, acSets)
         end
         # else return a normal term matcher
@@ -20,7 +21,7 @@ function matcher(val::Any, acSets)
 
     function literal_matcher(next, data, bindings)
         # car data is the first element of data
-        islist(data) && isequal(car(data), val) ? next(bindings, 1) : nothing
+        islist(data) && isequal(unwrap_const(car(data)), val) ? next(bindings, 1) : nothing
     end
 end
 
@@ -35,7 +36,7 @@ function matcher(slot::Slot, acSets)
                 return next(bindings, 1)
             end
         # elseif the first element of data matches the slot predicate, add it to bindings and call next
-        elseif slot.predicate(car(data))
+        elseif slot.predicate(unwrap_const(car(data)))
             rest = car(data)
             binds = assoc(bindings, slot.name, rest)
             next(binds, 1)
@@ -93,7 +94,7 @@ function matcher(segment::Segment, acSets)
             for i=length(data):-1:0
                 subexpr = take_n(data, i)
 
-                !segment.predicate(subexpr) && continue
+                !segment.predicate(unwrap_const(subexpr)) && continue
                 res = success(assoc(bindings, segment.name, subexpr), i)
                 res !== nothing && break
             end
@@ -104,7 +105,7 @@ function matcher(segment::Segment, acSets)
 end
 
 function term_matcher_constructor(term, acSets)
-    matchers = vcat([matcher(operation(term), acSets)], map(x -> matcher(x, acSets), parent(arguments(term))))
+    matchers = vcat([matcher(operation(term), acSets)], map(x -> matcher(unwrap_const(x), acSets), parent(arguments(term))))
 
     function loop(term, bindings′, matchers′) # Get it to compile faster
         if !islist(matchers′)
@@ -262,11 +263,11 @@ end
 # in the normal_matcher and in defslot_matcher and other_part_matcher
 function defslot_term_matcher_constructor(term, acSets)
     a = parent(arguments(term))
-    defslot_index = findfirst(x -> isa(x, DefSlot), a) # find the defslot in the term
-    defslot = a[defslot_index]
+    defslot_index = findfirst(x -> isa(unwrap_const(x), DefSlot), a) # find the defslot in the term
+    defslot = unwrap_const(a[defslot_index])
     defslot_matcher = matcher(defslot, acSets)
     if length(a) == 2
-        other_part_matcher = matcher(a[defslot_index == 1 ? 2 : 1], acSets)
+        other_part_matcher = matcher(unwrap_const(a[defslot_index == 1 ? 2 : 1]), acSets)
     else
         # if we hare here the operation is a multiplication or sum of n>2 terms
         # (because ^ cannot have more than 2 terms).

src/methods.jl

@@ -134,17 +134,17 @@ promote_symtype(::typeof(rem2pi), T::Type{<:Number}, mode) = T
 
 error_f_symbolic(f, T) = error("$f is not defined for $T.")
 
-function Base.rem2pi(x::Symbolic, mode::Base.RoundingMode)
+function Base.rem2pi(x::BasicSymbolic, mode::Base.RoundingMode)
     T = symtype(x)
     T <: Number ? term(rem2pi, x, mode) : error_f_symbolic(rem2pi, T)
 end
 
 # Specially handle inv and literal pow
-function Base.inv(x::Symbolic)
+function Base.inv(x::BasicSymbolic)
     T = symtype(x)
     T <: Number ? Base.:^(x, -1) : error_f_symbolic(rem2pi, T)
 end
-function Base.literal_pow(::typeof(^), x::Symbolic, ::Val{p}) where {p}
+function Base.literal_pow(::typeof(^), x::BasicSymbolic, ::Val{p}) where {p}
     T = symtype(x)
     T <: Number ? Base.:^(x, p) : error_f_symbolic(^, T)
 end
@@ -159,20 +159,20 @@ for f in monadic
     @eval promote_symtype(::$(typeof(f)), T::Type{<:LiteralRealImpl}) = LiteralReal
 end
 
-Base.:*(a::AbstractArray, b::Symbolic{<:Number}) = map(x->x*b, a)
-Base.:*(a::Symbolic{<:Number}, b::AbstractArray) = map(x->a*x, b)
+Base.:*(a::AbstractArray, b::BasicSymbolic{<:Number}) = map(x->x*b, a)
+Base.:*(a::BasicSymbolic{<:Number}, b::AbstractArray) = map(x->a*x, b)
 
 for f in [identity, one, zero, *, +, -]
     @eval promote_symtype(::$(typeof(f)), T::Type{<:Number}) = T
 end
 
 promote_symtype(::typeof(Base.real), T::Type{<:Number}) = Real
-Base.real(s::Symbolic{<:Number}) = islike(s, Real) ? s : term(real, s)
+Base.real(s::BasicSymbolic{<:Number}) = islike(s, Real) ? s : term(real, s)
 promote_symtype(::typeof(Base.conj), T::Type{<:Number}) = T
-Base.conj(s::Symbolic{<:Number}) = islike(s, Real) ? s : term(conj, s)
+Base.conj(s::BasicSymbolic{<:Number}) = islike(s, Real) ? s : term(conj, s)
 promote_symtype(::typeof(Base.imag), T::Type{<:Number}) = Real
-Base.imag(s::Symbolic{<:Number}) = islike(s, Real) ? zero(symtype(s)) : term(imag, s)
-Base.adjoint(s::Symbolic{<:Number}) = conj(s)
+Base.imag(s::BasicSymbolic{<:Number}) = islike(s, Real) ? zero(symtype(s)) : term(imag, s)
+Base.adjoint(s::BasicSymbolic{<:Number}) = conj(s)
 
 
 ## Booleans
@@ -186,29 +186,29 @@ for (f, Domain) in [(==) => Number, (!=) => Number,
                     xor => Bool]
     @eval begin
         promote_symtype(::$(typeof(f)), ::Type{<:$Domain}, ::Type{<:$Domain}) = Bool
-        (::$(typeof(f)))(a::Symbolic{<:$Domain}, b::$Domain) = term($f, a, b, type=Bool)
-        (::$(typeof(f)))(a::Symbolic{<:$Domain}, b::Symbolic{<:$Domain}) = term($f, a, b, type=Bool)
-        (::$(typeof(f)))(a::$Domain, b::Symbolic{<:$Domain}) = term($f, a, b, type=Bool)
+        (::$(typeof(f)))(a::BasicSymbolic{<:$Domain}, b::$Domain) = term($f, a, b, type=Bool)
+        (::$(typeof(f)))(a::BasicSymbolic{<:$Domain}, b::BasicSymbolic{<:$Domain}) = term($f, a, b, type=Bool)
+        (::$(typeof(f)))(a::$Domain, b::BasicSymbolic{<:$Domain}) = term($f, a, b, type=Bool)
     end
 end
 
 for f in [!, ~]
     @eval begin
         promote_symtype(::$(typeof(f)), ::Type{<:Bool}) = Bool
-        (::$(typeof(f)))(s::Symbolic{Bool}) = Term{Bool}(!, [s])
+        (::$(typeof(f)))(s::BasicSymbolic{Bool}) = Term{Bool}(!, [s])
     end
 end
 
 
 # An ifelse node
-function Base.ifelse(_if::Symbolic{Bool}, _then, _else)
+function Base.ifelse(_if::BasicSymbolic{Bool}, _then, _else)
     Term{Union{symtype(_then), symtype(_else)}}(ifelse, Any[_if, _then, _else])
 end
 promote_symtype(::typeof(ifelse), _, ::Type{T}, ::Type{S}) where {T,S} = Union{T, S}
 
 # Array-like operations
-Base.size(x::Symbolic{<:Number}) = ()
-Base.length(x::Symbolic{<:Number}) = 1
-Base.ndims(x::Symbolic{T}) where {T} = Base.ndims(T)
-Base.ndims(::Type{<:Symbolic{T}}) where {T} = Base.ndims(T)
-Base.broadcastable(x::Symbolic{T}) where {T<:Number} = Ref(x)
+Base.size(x::BasicSymbolic{<:Number}) = ()
+Base.length(x::BasicSymbolic{<:Number}) = 1
+Base.ndims(x::BasicSymbolic{T}) where {T} = Base.ndims(T)
+Base.ndims(::Type{<:BasicSymbolic{T}}) where {T} = Base.ndims(T)
+Base.broadcastable(x::BasicSymbolic{T}) where {T<:Number} = Ref(x)

src/ordering.jl

@@ -5,8 +5,8 @@
 <ₑ(a::Real,    b::Complex) = true
 <ₑ(a::Complex, b::Real) = false
 
-<ₑ(a::Symbolic, b::Number) = false
-<ₑ(a::Number,   b::Symbolic) = true
+<ₑ(a::BasicSymbolic, b::Number) = false
+<ₑ(a::Number, b::BasicSymbolic) = true
 
 <ₑ(a::Function, b::Function) = nameof(a) <ₑ nameof(b)
 
@@ -89,6 +89,8 @@ end
 
 function _get_degrees(::typeof(^), expr, degs_cache)
     base_expr, pow_expr = arguments(expr)
+    base_expr = unwrap_const(base_expr)
+    pow_expr = unwrap_const(pow_expr)
     if pow_expr isa Real
         @inbounds degs = map(_get_degrees(base_expr, degs_cache)) do (base, pow)
             (base => pow * pow_expr)
@@ -107,6 +109,8 @@ end
 
 function _get_degrees(::typeof(/), expr, degs_cache)
     nom_expr, denom_expr = arguments(expr)
+    nom_expr = unwrap_const(nom_expr)
+    denom_expr = unwrap_const(denom_expr)
     if denom_expr isa Number # constant denominator
         return _get_degrees(nom_expr, degs_cache)
     elseif nom_expr isa Number # constant nominator
@@ -157,11 +161,14 @@ function <ₑ(a::Tuple, b::Tuple)
 end
 
 function <ₑ(a::BasicSymbolic, b::BasicSymbolic)
+    if isconst(a) || isconst(b)
+        return <ₑ(unwrap_const(a), unwrap_const(b))
+    end
     da, db = get_degrees(a), get_degrees(b)
     fw = monomial_lt(da, db)
     bw = monomial_lt(db, da)
     if fw === bw && !isequal(a, b)
-        if _arglen(a) == _arglen(b)
+        if _arglen(a) == _arglen(b) != 0
             return (operation(a), arguments(a)...,) <ₑ (operation(b), arguments(b)...,)
         else
             return _arglen(a) < _arglen(b)