Merge pull request #723 from AayushSabharwal/as/small-array

feat: optimize small arrays in `BasicSymbolic`

Author: bowenszhu

docs/src/manual/rewrite.md

@@ -71,7 +71,7 @@ If you want to match a variable number of subexpressions at once, you will need
 @rule(+(~~xs) => ~~xs)(x + y + z)
 
 # output
-3-element view(::Vector{Any}, 1:3) with eltype Any:
+3-element view(::SymbolicUtils.SmallVec{Any, Vector{Any}}, 1:3) with eltype Any:
  z
  y
  x

src/SymbolicUtils.jl

@@ -25,6 +25,9 @@ import TaskLocalValues: TaskLocalValue
 
 include("cache.jl")
 Base.@deprecate istree iscall
+
+include("small_array.jl")
+
 export istree, operation, arguments, sorted_arguments, iscall
 # Sym, Term,
 # Add, Mul and Pow

src/small_array.jl

@@ -0,0 +1,133 @@
+"""
+    $(TYPEDEF)
+
+A mutable resizeable 3-length vector to implement non-allocating small vectors.
+
+$(TYPEDFIELDS)
+"""
+mutable struct Backing{T} <: AbstractVector{T}
+    """
+    Length of the buffer.
+    """
+    len::Int
+    x1::T
+    x2::T
+    x3::T
+
+    Backing{T}() where {T} = new{T}(0)
+    Backing{T}(x1) where {T} = new{T}(1, x1)
+    Backing{T}(x1, x2) where {T} = new{T}(2, x1, x2)
+    Backing{T}(x1, x2, x3) where {T} = new{T}(3, x1, x2, x3)
+end
+
+Base.size(x::Backing) = (x.len,)
+Base.isempty(x::Backing) = x.len == 0
+
+"""
+    $(TYPEDSIGNATURES)
+
+Value to use when removing an element from a `Backing`. Used so stored entries can be
+GC'ed when removed.
+"""
+defaultval(::Type{T}) where {T <: Number} = zero(T)
+defaultval(::Type{Any}) = nothing
+
+function Base.getindex(x::Backing, i::Int)
+    @boundscheck 1 <= i <= x.len
+    if i == 1
+        x.x1
+    elseif i == 2
+        x.x2
+    elseif i == 3
+        x.x3
+    end
+end
+
+function Base.setindex!(x::Backing, v, i::Int)
+    @boundscheck 1 <= i <= x.len
+    if i == 1
+        setfield!(x, :x1, v)
+    elseif i == 2
+        setfield!(x, :x2, v)
+    elseif i == 3
+        setfield!(x, :x3, v)
+    end
+end
+
+function Base.push!(x::Backing, v)
+    x.len < 3 || throw(ArgumentError("`Backing` is full"))
+    x.len += 1
+    x[x.len] = v
+end
+
+function Base.pop!(x::Backing{T}) where {T}
+    x.len > 0 || throw(ArgumentError("Array is empty"))
+    v = x[x.len]
+    x[x.len] = defaultval(T)
+    x.len -= 1
+    v
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Whether the `Backing` is full.
+"""
+isfull(x::Backing) = x.len == 3
+
+"""
+    $(TYPEDSIGNATURES)
+
+A small-buffer-optimized `AbstractVector`. Uses a `Backing` when the number of elements
+is within the size of `Backing`, and allocates a `V` when the number of elements exceed
+this limit.
+"""
+mutable struct SmallVec{T, V <: AbstractVector{T}} <: AbstractVector{T}
+    data::Union{Backing{T}, V}
+
+    function SmallVec{T}(x::AbstractVector{T}) where {T}
+        V = typeof(x)
+        if length(x) < 4
+            new{T, V}(Backing{T}(x...))
+        else
+            new{T, V}(x)
+        end
+    end
+
+    function SmallVec{T, V}() where {T, V}
+        new{T, V}(Backing{T}())
+    end
+
+    function SmallVec{T, V}(x::Union{Tuple, AbstractVector}) where {T, V}
+        if length(x) <= 3
+            new{T, V}(Backing{T}(x...))
+        else
+            new{T, V}(V(x isa Tuple ? collect(x) : x))
+        end
+    end
+end
+
+Base.convert(::Type{SmallVec{T, V}}, x::V) where {T, V} = SmallVec{T}(x)
+Base.convert(::Type{SmallVec{T, V}}, x) where {T, V} = SmallVec{T}(V(x))
+Base.convert(::Type{SmallVec{T, V}}, x::SmallVec{T, V}) where {T, V} = x
+
+Base.size(x::SmallVec) = size(x.data)
+Base.isempty(x::SmallVec) = isempty(x.data)
+Base.getindex(x::SmallVec, i::Int) = x.data[i]
+Base.setindex!(x::SmallVec, v, i::Int) = setindex!(x.data, v, i)
+
+function Base.push!(x::SmallVec{T, V}, v) where {T, V}
+    buf = x.data
+    buf isa Backing{T} || return push!(buf::V, v)
+    isfull(buf) || return push!(buf::Backing{T}, v)
+    x.data = V(buf)
+    return push!(x.data::V, v)
+end
+
+Base.pop!(x::SmallVec) = pop!(x.data)
+
+function Base.sizehint!(x::SmallVec{T, V}, n; kwargs...) where {T, V}
+    x.data isa Backing && return x
+    sizehint!(x.data, n; kwargs...)
+    x
+end

src/types.jl

@@ -16,7 +16,8 @@ const NO_METADATA = nothing
 sdict(kv...) = Dict{Any, Any}(kv...)
 
 using Base: RefValue
-const EMPTY_ARGS = []
+const SmallV{T} = SmallVec{T, Vector{T}}
+const EMPTY_ARGS = SmallV{Any}()
 const EMPTY_HASH = RefValue(UInt(0))
 const EMPTY_DICT = sdict()
 const EMPTY_DICT_T = typeof(EMPTY_DICT)
@@ -31,7 +32,7 @@ const ENABLE_HASHCONSING = Ref(true)
     end
     struct Term{T} <: BasicSymbolic{T}
         f::Any                 = identity  # base/num if Pow; issorted if Add/Dict
-        arguments::Vector{Any} = EMPTY_ARGS
+        arguments::SmallV{Any} = EMPTY_ARGS
         hash::RefValue{UInt}   = EMPTY_HASH
         hash2::RefValue{UInt} = EMPTY_HASH
     end
@@ -40,25 +41,25 @@ const ENABLE_HASHCONSING = Ref(true)
         dict::EMPTY_DICT_T     = EMPTY_DICT
         hash::RefValue{UInt}   = EMPTY_HASH
         hash2::RefValue{UInt} = EMPTY_HASH
-        arguments::Vector{Any} = EMPTY_ARGS
+        arguments::SmallV{Any} = EMPTY_ARGS
     end
     struct Add{T} <: BasicSymbolic{T}
         coeff::Any             = 0         # exp/den if Pow
         dict::EMPTY_DICT_T     = EMPTY_DICT
         hash::RefValue{UInt}   = EMPTY_HASH
         hash2::RefValue{UInt} = EMPTY_HASH
-        arguments::Vector{Any} = EMPTY_ARGS
+        arguments::SmallV{Any} = EMPTY_ARGS
     end
     struct Div{T} <: BasicSymbolic{T}
         num::Any               = 1
         den::Any               = 1
         simplified::Bool       = false
-        arguments::Vector{Any} = EMPTY_ARGS
+        arguments::SmallV{Any} = EMPTY_ARGS
     end
     struct Pow{T} <: BasicSymbolic{T}
         base::Any              = 1
         exp::Any               = 1
-        arguments::Vector{Any} = EMPTY_ARGS
+        arguments::SmallV{Any} = EMPTY_ARGS
     end
 end
 
@@ -564,17 +565,8 @@ function Sym{T}(name::Symbol; kw...) where {T}
     BasicSymbolic(s)
 end
 
-function unwrap_arr!(arr)
-    for i in eachindex(arr)
-        arr[i] = unwrap(arr[i])
-    end
-end
-
 function Term{T}(f, args; kw...) where T
-    if eltype(args) !== Any
-        args = convert(Vector{Any}, args)
-    end
-    unwrap_arr!(args)
+    args = SmallV{Any}(args)
 
     s = Term{T}(;f=f, arguments=args, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), kw...)
     BasicSymbolic(s)
@@ -606,7 +598,7 @@ function Add(::Type{T}, coeff, dict; metadata=NO_METADATA, kw...) where T
         end
     end
 
-    s = Add{T}(; coeff, dict, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), metadata, arguments=[], kw...)
+    s = Add{T}(; coeff, dict, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), metadata, arguments=SmallV{Any}(), kw...)
     BasicSymbolic(s)
 end
 
@@ -624,7 +616,7 @@ function Mul(T, a, b; metadata=NO_METADATA, kw...)
     else
         coeff = a
         dict = b
-        s = Mul{T}(; coeff, dict, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), metadata, arguments=[], kw...)
+        s = Mul{T}(; coeff, dict, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), metadata, arguments=SmallV{Any}(), kw...)
         BasicSymbolic(s)
     end
 end
@@ -645,7 +637,7 @@ ratio(x::Rat,y::Rat) = x//y
 function maybe_intcoeff(x)
     if ismul(x)
         if x.coeff isa Rational && isone(x.coeff.den)
-            Mul{symtype(x)}(; coeff=x.coeff.num, dict=x.dict, x.metadata, arguments=[])
+            Mul{symtype(x)}(; coeff=x.coeff.num, dict=x.dict, x.metadata, arguments=SmallV{Any}())
         else
             x
         end
@@ -692,7 +684,7 @@ function Div{T}(n, d, simplified=false; metadata=nothing, kwargs...) where {T}
         end
     end
 
-    s = Div{T}(; num=n, den=d, simplified, arguments=[], metadata)
+    s = Div{T}(; num=n, den=d, simplified, arguments=SmallV{Any}(), metadata)
     BasicSymbolic(s)
 end
 
@@ -712,7 +704,7 @@ function Pow{T}(a, b; metadata=NO_METADATA, kwargs...) where {T}
     b = unwrap(b)
     _iszero(b) && return 1
     _isone(b) && return a
-    s = Pow{T}(; base=a, exp=b, arguments=[], metadata)
+    s = Pow{T}(; base=a, exp=b, arguments=SmallV{Any}(), metadata)
     BasicSymbolic(s)
 end
 
@@ -728,13 +720,13 @@ function toterm(t::BasicSymbolic{T}) where T
         args = Any[]
         push!(args, t.coeff)
         for (k, coeff) in t.dict
-            push!(args, coeff == 1 ? k : Term{T}(E === MUL ? (^) : (*), Any[coeff, k]))
+            push!(args, coeff == 1 ? k : Term{T}(E === MUL ? (^) : (*), SmallV{Any}((coeff, k))))
         end
         Term{T}(operation(t), args)
     elseif E === DIV
-        Term{T}(/, Any[t.num, t.den])
+        Term{T}(/, SmallV{Any}((t.num, t.den)))
     elseif E === POW
-        Term{T}(^, [t.base, t.exp])
+        Term{T}(^, SmallV{Any}((t.base, t.exp)))
     else
         error_on_type()
     end
@@ -808,13 +800,13 @@ function makepow(a, b)
 end
 
 function term(f, args...; type = nothing)
-    args = map(unwrap, args)
+    args = SmallV{Any}(args)
     if type === nothing
         T = _promote_symtype(f, args)
     else
         T = type
     end
-    Term{T}(f, Any[args...])
+    Term{T}(f, args)
 end
 
 """
@@ -826,7 +818,7 @@ function unflatten(t::Symbolic{T}) where{T}
         f = operation(t)
         if f == (+) || f == (*)   # TODO check out for other n-ary --> binary ops
             a = arguments(t)
-            return foldl((x,y) -> Term{T}(f, Any[x, y]), a)
+            return foldl((x,y) -> Term{T}(f, SmallV{Any}((x, y))), a)
         end
     end
     return t
@@ -1192,7 +1184,7 @@ promote_symtype(f, Ts...) = Any
 
 struct FnType{X<:Tuple,Y,Z} end
 
-(f::Symbolic{<:FnType})(args...) = Term{promote_symtype(f, symtype.(args)...)}(f, Any[args...])
+(f::Symbolic{<:FnType})(args...) = Term{promote_symtype(f, symtype.(args)...)}(f, SmallV{Any}(args))
 
 function (f::Symbolic)(args...)
     error("Sym $f is not callable. " *