io_delimited.cr

class IO::Delimited
  include IO

  def self.new(io : IO, delimiter : String)
    new(io, delimiter.to_slice)
  end

  def initialize(@io : IO, @delimiter : Slice(UInt8))
    raise ArgumentError.new("empty delimiter") if @delimiter.size == 0
    @reached_end = false
  end

  def read(slice : Slice(UInt8))
    return 0 if @reached_end

    count = slice.size
    return 0 if count == 0

    total_read = 0

    while count > 0
      byte = internal_read_byte
      unless byte
        raise "couldn't find delimiter"
      end

      # If we found the first byte of the delimiter
      if byte == @delimiter.to_unsafe[0]
        matched_delimiter = true

        # We try to match the other bytes of the delimiter
        (@delimiter.size - 1).times do |i|
          byte = internal_read_byte
          unless byte
            raise "couldn't find delimiter"
          end

          # If the next byte doesn't match the next byte in the delimiter
          unless byte == @delimiter.to_unsafe[i + 1]
            # It might happen that the delimiter can include a prefix
            # of itself, so we can't just discard all previous read bytes.
            # What we do is we just discard the first byte,
            # writing it to the target slice
            slice.to_unsafe.value = @delimiter.to_unsafe[0]
            slice += 1
            total_read += 1
            count -= 1

            # And we write the rest to the internal buffer, making
            # sure not to move the current pointer inside the buffer (pos)
            buffer = @buffer ||= MemoryIO.new
            pos = buffer.pos
            buffer.write(Slice.new(@delimiter.to_unsafe + 1, i))
            buffer.write_byte(byte)
            buffer.pos = pos
            matched_delimiter = false
            break
          end
        end

        # We found the delimiter
        if matched_delimiter
          @reached_end = true
          break
        end
      else
        # Otherwise, just read into the slice
        slice.to_unsafe.value = byte
        slice += 1
        total_read += 1
        count -= 1
      end
    end

    total_read
  end

  def write(slice : Slice(UInt8))
    raise "read-only stream"
  end

  # We read from the intermediate buffer, if it has bytes,
  # otherwise from the wrapped IO
  private def internal_read_byte
    if buffer = @buffer
      byte = buffer.read_byte
      unless byte
        buffer.clear
        byte = @io.read_byte
      end
    else
      byte = @io.read_byte
    end
    byte
  end
end

mem_io = MemoryIO.new("hello world, lolonlolololng this is a lolong string")
io = IO::Delimited.new(mem_io, "lolong")
p io.gets_to_end
p mem_io.gets_to_end

mem_io.rewind
io = IO::Delimited.new(mem_io, "w")
p io.gets_to_end
p mem_io.gets_to_end

mem_io.rewind
io = IO::Delimited.new(mem_io, "long")
p io.gets_to_end
p mem_io.gets_to_end