tap_stdout.cr

private def overwrite_memory(src : T, dst : T) : Nil forall T
  src = src.as(T*).as(UInt8*)
  dst = dst.as(T*).as(UInt8*)
  src.copy_to(dst, instance_sizeof(T))
end

abstract class Capture
  abstract def transform(bytes : Bytes) : Bytes

  @done = Channel(Nil).new
  @stream : IO::FileDescriptor? = nil
  @original_stream : IO::FileDescriptor? = nil

  def tap_output_stream(stream : IO::FileDescriptor) : Nil
    raise "Already captured." if @stream

    @stream = stream
    @original_stream = original_stream = stream.dup
    reader, writer = IO.pipe
    overwrite_memory(src: writer, dst: stream)

    spawn do
      buf = Bytes.new(original_stream.buffer_size)
      loop do
        size = reader.read(buf)
        break if size == 0
        original_stream.write(transform(buf[0, size]))
      end
      original_stream.flush
      @done.send(nil)
    end
  end

  def untap : Nil
    if (stream = @stream)
      stream.close
      @done.receive
      if (original_stream = @original_stream)
        overwrite_memory(src: original_stream, dst: stream)
        @original_stream = nil
      end
      @stream = nil
    end
  end

  def self.tap_output_stream(stream : IO::FileDescriptor) : Nil
    capture = new
    capture.tap_output_stream(stream)
    begin
      yield
    ensure
      capture.untap
    end
  end
end

class Rot13Capture < Capture
  def transform(buf : Bytes) : Bytes
    buf.map! do |c|
      case c
      when 'a'.ord..'z'.ord
        (c - 'a'.ord + 13) % 26 + 'a'.ord
      when 'A'.ord..'Z'.ord
        (c - 'A'.ord + 13) % 26 + 'A'.ord
      else c
      end
    end
  end
end

puts "Before"
Rot13Capture.tap_output_stream(STDOUT) do
  puts "Hello, World!"
end
puts "After"