Allow concatenating small InlineStrings to return an InlineString (#59)

* Allow concatenating small InlineStrings to return an InlineString

* More specific test for issue#2

* fixup! More specific test for issue#2

* Update src/InlineStrings.jl

* fixup! Update src/InlineStrings.jl

* Add more string concat tests

* Add comments

* fixup! Add more string concat tests

* More micro-optimizations for tiny inline strings

* fixup! More micro-optimizations for tiny inline strings

Author: nickrobinson251

src/InlineStrings.jl

@@ -73,6 +73,8 @@ for sz in (1, 4, 8, 16, 32, 64, 128, 256)
     end
 end
 
+const SmallInlineStrings = Union{String1, String3, String7, String15}
+
 # used to zero out n lower bytes of an inline string
 clear_n_bytes(s, n) = Base.shl_int(Base.lshr_int(s, 8 * n), 8 * n)
 _bswap(x::T) where {T <: InlineString} = T === InlineString1 ? x : Base.bswap_int(x)
@@ -595,10 +597,11 @@ end
 end
 
 const BaseStrs = Union{Char, String, SubString{String}}
-Base.string(a::InlineString) = String(a)
+Base.string(a::InlineString) = a
 Base.string(a::InlineString...) = _string(a...)
 Base.string(a::BaseStrs, b::InlineString) = _string(a, b)
 Base.string(a::BaseStrs, b::BaseStrs, c::InlineString) = _string(a, b, c)
+
 @inline function _string(a::Union{BaseStrs, InlineString}...)
     n = 0
     for v in a
@@ -616,6 +619,40 @@ Base.string(a::BaseStrs, b::BaseStrs, c::InlineString) = _string(a, b, c)
     return out
 end
 
+# For more and/or bigger InlineStrings creating a `Base.String` is faster
+const _SmallerInlineStrings = Union{InlineString1, InlineString3, InlineString7}
+Base.string(a::_SmallerInlineStrings, b::_SmallerInlineStrings, c::_SmallerInlineStrings) =
+    _string(_string(a, b), c)
+const _SmallestInlineStrings = Union{InlineString1, InlineString3}
+Base.string(a::_SmallestInlineStrings, b::_SmallestInlineStrings, c::_SmallestInlineStrings, d::_SmallestInlineStrings) =
+    _string(_string(_string(a, b), c), d)
+
+# Only benefit from keeping the small-ish strings as InlineStrings
+function _string(a::Ta, b::Tb) where {Ta <: SmallInlineStrings, Tb <: SmallInlineStrings}
+    T = summed_type(Ta, Tb)
+    lb_a = Int(!isa(a, InlineString1)) # no "length byte" to remove if InlineString1
+    lb_b = Int(!isa(b, InlineString1))
+    len_a = sizeof(a)
+    len_b = sizeof(b)
+    # Remove length byte (lshr), grow to new size (zext), move chars forward (shl).
+    a2 = Base.shl_int(Base.zext_int(T, Base.lshr_int(a, 8*lb_a)), 8 * (sizeof(T) - sizeof(Ta) + lb_a))
+    b2 = Base.shl_int(Base.zext_int(T, Base.lshr_int(b, 8*lb_b)), 8 * (sizeof(T) - sizeof(Tb) + lb_b - len_a))
+    lb = _oftype(T, len_a + len_b)  # new length byte
+    return Base.or_int(Base.or_int(a2, b2), lb)
+end
+
+summed_type(::Type{InlineString1}, ::Type{InlineString1}) = InlineString3
+summed_type(::Type{InlineString3}, ::Type{InlineString1}) = InlineString7
+summed_type(::Type{InlineString3}, ::Type{InlineString3}) = InlineString7
+summed_type(::Type{InlineString7}, ::Type{InlineString1}) = InlineString15
+summed_type(::Type{InlineString7}, ::Type{InlineString3}) = InlineString15
+summed_type(::Type{InlineString7}, ::Type{InlineString7}) = InlineString15
+summed_type(::Type{InlineString15}, ::Type{InlineString1}) = InlineString31
+summed_type(::Type{InlineString15}, ::Type{InlineString3}) = InlineString31
+summed_type(::Type{InlineString15}, ::Type{InlineString7}) = InlineString31
+summed_type(::Type{InlineString15}, ::Type{InlineString15}) = InlineString31
+summed_type(a::Type{<:SmallInlineStrings}, b::Type{<:SmallInlineStrings}) = summed_type(b, a)
+
 function Base.repeat(x::T, r::Integer) where {T <: InlineString}
     r < 0 && throw(ArgumentError("can't repeat a string $r times"))
     r == 0 && return ""
@@ -886,9 +923,6 @@ end
 ## InlineString sorting
 using Base.Sort, Base.Order
 
-# Only small-ish InlineStrings benefit from RadixSort algorithm
-const SmallInlineStrings = Union{String1, String3, String7, String15}
-
 # And under certain thresholds, MergeSort is faster than RadixSort, even for small InlineStrings
 const MergeSortThresholds = Dict(
     1 => 2^5,
@@ -900,6 +934,7 @@ const MergeSortThresholds = Dict(
 struct InlineStringSortAlg <: Algorithm end
 const InlineStringSort = InlineStringSortAlg()
 
+# Only small-ish InlineStrings benefit from RadixSort algorithm
 Base.Sort.defalg(::AbstractArray{<:Union{SmallInlineStrings, Missing}}) = InlineStringSort
 
 struct Radix

test/runtests.jl

@@ -69,8 +69,7 @@ x = InlineString1(buf)
 @test InlineString(String1("a")) === String1("a")
 
 # https://github.com/JuliaData/InlineStrings.jl/issues/2
-x = InlineString("hey")
-@test typeof(string(x)) == String
+@test eltype(string.(AbstractString[])) == AbstractString
 
 # https://github.com/JuliaData/InlineStrings.jl/issues/8
 # construction from pointer
@@ -277,6 +276,7 @@ const INLINES = map(InlineString, STRINGS)
         @test Array{UInt8}(x) == Array{UInt8}(y)
         @test isascii(x) == isascii(y)
         @test x * x == y * y
+        @test x * x * x == y * y * y
         @test x^5 == y^5
         @test string(x) == string(y)
         @test join([x, x]) == join([y, y])
@@ -349,6 +349,28 @@ const INLINES = map(InlineString, STRINGS)
     end
 end
 
+@testset "`string` / `*`" begin
+    # Check `string` overload handles `String1` being concat with other small InlineStrings,
+    # because it is easy to mishandle `String1` as it doesn't have a length byte.
+    a = "a"
+    @test String1(a) * String1(a) == a * a
+    @test String1(a) * String1(a) isa InlineString3
+    b = "bb"
+    @test String1(a) * String3(b) == a * b
+    @test String1(a) * String3(b) isa InlineString7
+    @test String1(a) * String7(b) == a * b
+    @test String1(a) * String7(b) isa InlineString15
+    @test String1(a) * String15(b) == a * b
+    @test String1(a) * String15(b) isa InlineString31
+    @test String1(a) * String3(b) * String7(b) == a * b * b
+    @test String1(a) * String3(b) * String7(b) isa InlineString15
+    # Check some other combination of small inline strings also work as expected
+    @test String3(a) * String7(b) == a * b
+    @test String3(a) * String7(b) isa InlineString15
+    @test String3(a) * String3(b) * String7(b) == a * b * b
+    @test String3(a) * String3(b) * String7(b) isa InlineString15
+end
+
 @testset "InlineString parsing" begin
 testcases = [
     ("", InlineString7(""), NamedTuple(), OK | EOF),