Improve the annotated join method and dispatch (#51914)

With the initial implementation, join could work for AnnotatedStrings,
however only when the eltype of the iterator or delim was itself a
Annotated{String,Char}. This was better than nothing, but seems
inconsistent in the face of mixed iterators.

Furthermore, the implementation of an annotated join was far from
optimised, relying on zipping and then calling annotatedstring(xs...).
By contrast, the non-annotated implementation relies on printing to IO
and even has manually defined alternative methods for optional arguments
to minimise code generation.

With the advent of AnnotatedIOBuffer and _isannotated, we can improve on
both those aspects. The new AnnotatedIOBuffer type allows us to re-use
the optimised join(::IO, ...) methods, and we can more reliably dispatch
to them with _isannotated. Since this is a type-based decision, the
Julia compiler is kind enough to work out which branch is taken at
compile-time, making this zero-overhead in the unannotated case.

Author: tecosaur

base/strings/io.jl

@@ -353,28 +353,19 @@ function join(io::IO, iterator, delim="")
     end
 end
 
-# TODO: If/when we have `AnnotatedIO`, we can revisit this and
-# implement it more nicely.
-function join_annotated(iterator, delim="", last=delim)
-    xs = zip(iterator, Iterators.repeated(delim)) |> Iterators.flatten |> collect
-    xs = xs[1:end-1]
-    if length(xs) > 1
-        xs[end-1] = last
-    end
-    annotatedstring(xs...)::AnnotatedString{String}
-end
-
-function _join_maybe_annotated(args...)
-    if any(_isannotated ∘ eltype, args)
-        join_annotated(args...)
+function _join_preserve_annotations(iterator, args...)
+    if _isannotated(eltype(iterator)) || any(_isannotated, args)
+        io = AnnotatedIOBuffer()
+        join(io, iterator, args...)
+        read(seekstart(io), AnnotatedString{String})
     else
-        sprint(join, args...)
+        sprint(join, iterator, args...)
     end
 end
 
-join(iterator) = _join_maybe_annotated(iterator)
-join(iterator, delim) = _join_maybe_annotated(iterator, delim)
-join(iterator, delim, last) = _join_maybe_annotated(iterator, delim, last)
+join(iterator) = _join_preserve_annotations(iterator)
+join(iterator, delim) = _join_preserve_annotations(iterator, delim)
+join(iterator, delim, last) = _join_preserve_annotations(iterator, delim, last)
 
 ## string escaping & unescaping ##
 

test/strings/annotated.jl

@@ -91,7 +91,7 @@ end
                                       (6:11, :other => 0x02)])
     str1 = Base.AnnotatedString("test", [(1:4, :label => 5)])
     str2 = Base.AnnotatedString("case", [(2:3, :label => "oomph")])
-    @test join([str1, str1], Base.AnnotatedString(" ")) ==
+    @test join([str1, str1], ' ') ==
         Base.AnnotatedString("test test",
                      [(1:4, :label => 5),
                       (6:9, :label => 5)])