Insert bytes directly into the stream

This patch removes the need for backing up bytes in every function by
simply splicing the bytes into place directly.

src/Runic.jl

@@ -87,11 +87,10 @@ end
 
 # Write formatted thing and reset the output stream
 function write_and_reset(ctx::Context, bytes::Union{String, AbstractVector{UInt8}})
-    fmt_pos = position(ctx.fmt_io)
-    nb = write(ctx.fmt_io, bytes)
-    seek(ctx.fmt_io, fmt_pos)
-    @assert nb == (bytes isa Vector{UInt8} ? length(bytes) : sizeof(bytes))
-    return nb
+    return write_and_reset(ctx.fmt_io, bytes)
+end
+function insert_bytes(ctx::Context, bytes::Union{String, AbstractVector{UInt8}}, sz::Integer)
+    return insert_bytes(ctx.fmt_io, bytes, Int(sz))
 end
 
 struct NullNode end
@@ -109,11 +108,6 @@ function format_node_with_children!(ctx::Context, node::JuliaSyntax.GreenNode)
     ctx.prev_sibling = nothing
     ctx.next_sibling = nothing
 
-    # A formatted node can have a larger span than the original so we need to backup the
-    # original bytes and keep track of the accumulated span of processed children
-    original_bytes = node_bytes(ctx, node) # TODO: Read into reusable buffer?
-    span_sum = 0
-
     # The new node parts. `children′` aliases `children` and only copied below if any of the
     # nodes change ("copy-on-write").
     head′ = JuliaSyntax.head(node)
@@ -127,7 +121,6 @@ function format_node_with_children!(ctx::Context, node::JuliaSyntax.GreenNode)
         ctx.prev_sibling = get(children′, i - 1, nothing)
         ctx.next_sibling = get(children, i + 1, nothing)
         child′ = child
-        span_sum += JuliaSyntax.span(child)
         this_child_changed = false
         itr = 0
         # Loop until this node reaches a steady state and is accepted
@@ -158,13 +151,6 @@ function format_node_with_children!(ctx::Context, node::JuliaSyntax.GreenNode)
                 error("infinite loop?")
             end
         end
-        if this_child_changed
-            # If the node changed we have to re-write the original bytes for the next
-            # children to the output stream and then reset
-            remaining_bytes = @view original_bytes[(span_sum + 1):end]
-            nb = write_and_reset(ctx, remaining_bytes)
-            @assert nb == length(remaining_bytes)
-        end
         any_child_changed |= this_child_changed
         if any_child_changed
             # De-alias the children if not already done

src/chisels.jl

@@ -123,3 +123,47 @@ function first_non_whitespace_child(node::JuliaSyntax.GreenNode)
     idx = findfirst(!JuliaSyntax.is_whitespace, children)::Int
     return children[idx]
 end
+
+##########################
+# Utilities for IOBuffer #
+##########################
+
+function write_and_reset(io::IOBuffer, bytes::Union{String, AbstractVector{UInt8}})
+    pos = position(io)
+    nb = write(io, bytes)
+    @assert nb == (bytes isa Vector{UInt8} ? length(bytes) : sizeof(bytes))
+    seek(io, pos)
+    @assert pos == position(io)
+    return nb
+end
+
+# Insert bytes at the current position in the IOBuffer. `size` is the number of bytes until
+# the next node starts (i.e., `size` is the span of the node we are replacing). If `size` is
+# smaller than the length of `bytes` need to shift all remainig bytes in the buffer to the
+# right to make room, otherwise we shift the remaining bytes to the left to fill the gap.
+function insert_bytes(io::IOBuffer, bytes::Union{String, AbstractVector{UInt8}}, size::Int)
+    pos = position(io)
+    nb = (bytes isa AbstractVector{UInt8} ? length(bytes) : sizeof(bytes))
+    if nb == size
+        nw = write(io, bytes)
+        @assert nb == nw
+    else
+        backup = IOBuffer() # TODO: global const with lock?
+        seek(io, pos + size)
+        @assert position(io) == pos + size
+        nb_written_to_backup = write(backup, io)
+        seek(io, pos)
+        @assert position(io) == pos
+        nw = write(io, bytes)
+        @assert nb == nw
+        nb_read_from_backup = write(io, seekstart(backup))
+        @assert nb_written_to_backup == nb_read_from_backup
+        truncate(io, position(io))
+    end
+    seek(io, pos)
+    @assert position(io) == pos
+    return nb
+end
+
+insert_bytes(io::IOBuffer, bytes::Union{String, AbstractVector{UInt8}}, size::Integer) =
+    insert_bytes(io, bytes, Int(size))

src/runestone.jl

@@ -18,7 +18,7 @@ function trim_trailing_whitespace(ctx::Context, node::JuliaSyntax.GreenNode)
         return nothing
     end
     # Write new bytes and reset the stream
-    nb = write_and_reset(ctx, str′)
+    nb = insert_bytes(ctx, str′, JuliaSyntax.span(node))
     @assert nb != JuliaSyntax.span(node)
     # Create new node and return it
     node′ = JuliaSyntax.GreenNode(JuliaSyntax.head(node), nb, ())
@@ -43,7 +43,7 @@ function format_hex_literals(ctx::Context, node::JuliaSyntax.GreenNode)
     while length(bytes) < target_spans[i]
         insert!(bytes, 3, '0')
     end
-    nb = write_and_reset(ctx, bytes)
+    nb = insert_bytes(ctx, bytes, span)
     @assert nb == length(bytes) == target_spans[i]
     # Create new node and return it
     node′ = JuliaSyntax.GreenNode(JuliaSyntax.head(node), nb, ())
@@ -81,7 +81,7 @@ function format_oct_literals(ctx::Context, node::JuliaSyntax.GreenNode)
     while length(bytes) < target_span
         insert!(bytes, 3, '0')
     end
-    nb = write_and_reset(ctx, bytes)
+    nb = insert_bytes(ctx, bytes, span)
     @assert nb == length(bytes) == target_span
     # Create new node and return it
     node′ = JuliaSyntax.GreenNode(JuliaSyntax.head(node), nb, ())
@@ -155,8 +155,6 @@ function spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, is_x::F) whe
     children = verified_children(node)
     children′ = children
     any_changes = false
-    original_bytes = node_bytes(ctx, node)
-    span_sum = 0
     pos = position(ctx.fmt_io)
     ws = JuliaSyntax.GreenNode(
         JuliaSyntax.SyntaxHead(K"Whitespace", JuliaSyntax.TRIVIA_FLAG), 1, (),
@@ -167,7 +165,6 @@ function spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, is_x::F) whe
     looking_for_x = false
 
     for (i, child) in pairs(children)
-        span_sum += JuliaSyntax.span(child)
         if JuliaSyntax.kind(child) === K"NewlineWs" ||
             (i == 1 && JuliaSyntax.kind(child) === K"Whitespace")
             # NewlineWs are accepted as is by this pass.
@@ -190,11 +187,8 @@ function spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, is_x::F) whe
                     children′ = children[1:i - 1]
                 end
                 push!(children′, ws)
-                write_and_reset(ctx, " ")
+                insert_bytes(ctx, " ", JuliaSyntax.span(child))
                 accept_node!(ctx, ws)
-                # Re-write bytes for remaining children
-                remaining_bytes = @view original_bytes[(span_sum + 1):end]
-                write_and_reset(ctx, remaining_bytes)
                 looking_for_whitespace = false
             elseif JuliaSyntax.haschildren(child) &&
                     JuliaSyntax.kind(first_leaf(child)) === K"Whitespace"
@@ -213,9 +207,7 @@ function spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, is_x::F) whe
                     @assert JuliaSyntax.span(child′) == JuliaSyntax.span(child) - JuliaSyntax.span(child_ws) + 1
                     bytes_to_skip = JuliaSyntax.span(child) - JuliaSyntax.span(child′)
                     @assert bytes_to_skip > 0
-                    remaining_bytes_inclusive =
-                        @view original_bytes[(span_sum + 1 + bytes_to_skip - JuliaSyntax.span(child)):end]
-                    write_and_reset(ctx, remaining_bytes_inclusive)
+                    insert_bytes(ctx, "", bytes_to_skip)
                     accept_node!(ctx, child′)
                     any_changes = true
                     if children′ === children
@@ -238,13 +230,10 @@ function spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, is_x::F) whe
                     children′ = children[1:i - 1]
                 end
                 push!(children′, ws)
-                write_and_reset(ctx, " ")
+                insert_bytes(ctx, " ", 0)
                 accept_node!(ctx, ws)
                 # Write and accept the node
                 push!(children′, child)
-                remaining_bytes_inclusive =
-                    @view original_bytes[(span_sum + 1 - JuliaSyntax.span(child)):end]
-                write_and_reset(ctx, remaining_bytes_inclusive)
                 accept_node!(ctx, child)
                 looking_for_whitespace = JuliaSyntax.kind(last_leaf(child)) !== K"Whitespace"
                 if looking_for_x
@@ -313,8 +302,6 @@ function no_spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, is_x::F)
     children = verified_children(node)
     children′ = children
     any_changes = false
-    original_bytes = node_bytes(ctx, node)
-    span_sum = 0
     pos = position(ctx.fmt_io)
 
     looking_for_x = false
@@ -326,7 +313,6 @@ function no_spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, is_x::F)
     end
 
     for (i, child) in pairs(children)
-        span_sum += JuliaSyntax.span(child)
         if (i == 1 || i == length(children)) && JuliaSyntax.kind(child) === K"Whitespace"
             accept_node!(ctx, child)
             any_changes && push!(children′, child)
@@ -336,8 +322,7 @@ function no_spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, is_x::F)
             if children′ === children
                 children′ = children[1:i - 1]
             end
-            remaining_bytes = @view original_bytes[(span_sum + 1):end]
-            write_and_reset(ctx, remaining_bytes)
+            insert_bytes(ctx, "", JuliaSyntax.span(child))
         else
             @assert JuliaSyntax.kind(child) !== K"Whitespace"
             if looking_for_x
@@ -394,24 +379,18 @@ function replace_with_in(ctx::Context, node::JuliaSyntax.GreenNode)
     @assert JuliaSyntax.kind(in_node) in KSet"∈ ="
     @assert JuliaSyntax.is_trivia(in_node)
     @assert is_leaf(in_node)
-    bytes = node_bytes(ctx, node) # TODO: Need something better...
     # Accept nodes to advance the stream
-    span_sum = 0
     for i in 1:(in_index - 1)
-        span_sum += JuliaSyntax.span(children[i])
         accept_node!(ctx, children[i])
     end
-    span_sum += JuliaSyntax.span(children[in_index])
     # Construct the replacement
-    nb = write_and_reset(ctx, "in")
+    nb = insert_bytes(ctx, "in", JuliaSyntax.span(in_node))
     in_node′ = JuliaSyntax.GreenNode(
         JuliaSyntax.SyntaxHead(K"in", JuliaSyntax.TRIVIA_FLAG), nb, (),
     )
     accept_node!(ctx, in_node′)
     children′ = copy(children)
     children′[in_index] = in_node′
-    # Write the backed up bytes
-    write_and_reset(ctx, @view(bytes[(span_sum + 1):end]))
     # Accept remaining eq_children
     for i in (in_index + 1):length(children′)
         accept_node!(ctx, children′[i])
@@ -424,19 +403,14 @@ function replace_with_in_cartesian(ctx::Context, node::JuliaSyntax.GreenNode)
     @assert JuliaSyntax.kind(node) === K"cartesian_iterator" && !is_leaf(node)
     children = verified_children(node)
     children′ = children
-    bytes = node_bytes(ctx, node)
-    span_sum = 0
     for (i, child) in pairs(children)
-        span_sum += JuliaSyntax.span(child)
         if JuliaSyntax.kind(child) === K"="
             child′ = replace_with_in(ctx, child)
             if child′ !== nothing
                 if children′ === children
                     children′ = copy(children)
                 end
                 children′[i] = child′
-                # Need to re-write the bytes
-                write_and_reset(ctx, @view(bytes[(span_sum + 1):end]))
             else
                 children′[i] = child
                 accept_node!(ctx, child)
@@ -463,8 +437,7 @@ function for_loop_use_in(ctx::Context, node::JuliaSyntax.GreenNode)
     )
         return nothing
     end
-    pos = position(ctx.fmt_io) # In case a reset is needed later
-    bytes = node_bytes(ctx, node)
+    pos = position(ctx.fmt_io)
     children = verified_children(node)
     for_index = findfirst(c -> JuliaSyntax.kind(c) === K"for" && is_leaf(c), children)::Int
     for_node = children[for_index]
@@ -484,17 +457,13 @@ function for_loop_use_in(ctx::Context, node::JuliaSyntax.GreenNode)
         for_spec_node′ = replace_with_in_cartesian(ctx, for_spec_node)
     end
     if for_spec_node′ === nothing
-        # TODO: reset here? Because we werent' supposed to write "for" above
         seek(ctx.fmt_io, pos)
         return nothing
     end
     @assert position(ctx.fmt_io) == pos + mapreduce(JuliaSyntax.span, +, @view(children[1:for_index])) + JuliaSyntax.span(for_spec_node′)
     # Insert the new for spec node
     children′ = copy(children)
     children′[for_spec_index] = for_spec_node′
-    # Write the backup bytes
-    span_sum = JuliaSyntax.span(for_node) + JuliaSyntax.span(for_spec_node)
-    write_and_reset(ctx, @view(bytes[(span_sum + 1):end]))
     # At this point the eq node is done, just accept any remaining nodes
     # TODO: Don't need to do this...
     for i in (for_spec_index + 1):length(children′)

test/runtests.jl

@@ -11,6 +11,31 @@ using JuliaSyntax:
     # Type stability of verified_children
     node = JuliaSyntax.parseall(JuliaSyntax.GreenNode, "a = 1 + b\n")
     @test typeof(@inferred Runic.verified_children(node)) <: Vector{<:JuliaSyntax.GreenNode}
+
+    # insert_bytes: insert larger
+    io = IOBuffer(); write(io, "abc"); seek(io, 1)
+    p = position(io)
+    Runic.insert_bytes(io, "xx", 1)
+    @test p == position(io)
+    @test read(io, String) == "xxc"
+    seekstart(io)
+    @test read(io, String) == "axxc"
+    # insert_bytes: insert smaller
+    io = IOBuffer(); write(io, "abbc"); seek(io, 1)
+    p = position(io)
+    Runic.insert_bytes(io, "x", 2)
+    @test p == position(io)
+    @test read(io, String) == "xc"
+    seekstart(io)
+    @test read(io, String) == "axc"
+    # insert_bytes: insert same
+    io = IOBuffer(); write(io, "abc"); seek(io, 1)
+    p = position(io)
+    Runic.insert_bytes(io, "x", 1)
+    @test p == position(io)
+    @test read(io, String) == "xc"
+    seekstart(io)
+    @test read(io, String) == "axc"
 end
 
 @testset "Trailing whitespace" begin