My solution to Puzzle Node's "Hitting Rock Bottom"

rockbottom.rb

# TODO
#
# Handle adding more water units than cave capacity (use something like
# cave.overflow?)

module RockBottom
  #
  # Represents one column in the cave.
  #
  # You only care about how many water units the column contains - which is
  # why we only keep track of it's size and not the actual '~'
  #
  # Instance variables:
  # - @max_size : maximum number of water units allowed
  # - @bottom :  deepest possible depth at which a water unit can be placed
  # - @size : number of water units stored
  # - @cascades : number of water units that are currently falling down from
  #   some point higher than @bottom, without actually having reached it yet
  #
  # Allowed operations are:
  # - flow : Add one water unit to the bottom.
  # - cascade : Add one water unit to the top.
  #
  class Column
    attr_reader :max_size, :size, :bottom

    def initialize(max_size)
      @max_size = max_size
      @size = 0
      @cascades = 0
      @bottom = @max_size - 1
    end

    # Add one water unit at @bottom
    def flow
      if @size < @max_size
        @size += 1
        @bottom -= 1
      end
    end

    #
    # Add one water unit at <height>.
    #
    # If there were other units that have cascaded down before this one, push
    # them down. If the stacked cascading units reach a height of <height>,
    # then they have touched the surface of the cave and no more cascading is
    # allowed.
    #
    # The size of the column can't be measured while cascading is in process.
    # Only columns that have water units on the cave surface are considered to
    # have a size.
    #
    # Returns a boolean indicating if the cascading process finished.
    #
    def cascade(height)
      @cascades += 1
      @bottom = height - 1 if @cascades == 1
      touched_bottom = false
      if @cascades == @max_size - height
        @cascades = 0
        @size = @max_size - height
        touched_bottom = true       
      else
        @size = nil
      end
      touched_bottom
    end

  end

  #
  # Represents a cave with irregular floor and water units streaming in.
  #
  # Water can either be :cascading or :flowing.
  #
  # Cascading happens when the tip of the flow of water is at a higher level
  # than the cave surface of the corresponding column, so water has to fall
  # down until it reaches that point.
  #
  # Flowing happens when the waterflow has touched the surface of one column,
  # so it has to flow onto the one to it's right.
  #
  # Instance methods:
  # - pump : add a certain amount of water units to the cave.
  # - map : returns a string representation of the cave.
  # - heights : returns and Array containing the number of water units in every column
  #
  class Cave
    attr_reader :columns

    def initialize(src)
      @columns = build_columns(src)
      @current_depth = 0
      @cascading_height = 0
      @next_column_index = 0
      @influx_state = :flowing
    end

    def map
      s = "\t" + ("#" * (@columns.count + 1))
      @map.each_with_index { |l, i| s += (i.to_s + "\t" + l + "#") }
      s += ("\t" + ("#" * (@columns.count + 1)))
    end

    def heights
      @columns.map { |col| col.size } + [0]
    end

    #
    # Add <units> units of water to the cave; update the map accordingly.
    #
    # Takes a block that receives the heights and map after every iteration;
    # use for debugging.
    #
    def pump(units=1)
      units.times do |t|
        column = @columns[@next_column_index]
        if @influx_state == :flowing
          flow column
        elsif @influx_state == :cascading
          cascade column
        end
        find_next_column
        if block_given?
          yield heights, map
        end       
      end
    end

    private

    def flow(column)
      column.flow
      mark_flow
    end

    def cascade(column)
      touched_bottom = column.cascade(@cascading_height)
      mark_cascade
      if touched_bottom
        @current_depth = @cascading_height + column.size - 1
        @influx_state = :flowing
        @cascading_height = 0
      end
    end

    def mark_flow
      @map[@current_depth][@next_column_index] = '~'
    end

    def mark_cascade
      tmp = @cascading_height
      while @map[tmp][@next_column_index] == '~'
        tmp += 1
      end
      @map[tmp][@next_column_index] = '~'
    end

    #
    # Returns the column to add the next water unit to.
    #
    # If water flow is cascading, keep adding units to that column until they
    # touch the bottom, at which point water starts flowing. When flowing,
    # always try to go right, but if the water hits a wall, start filling the
    # level above from the left.
    #
    def find_next_column
      return if @influx_state == :cascading
      right_column = @columns[@next_column_index + 1]
      if right_column.bottom >= @current_depth
        next_column_right
      else
        next_column_left
      end
    end

    def next_column_right
      @next_column_index += 1
      if @columns[@next_column_index].bottom > @current_depth
        @influx_state = :cascading
        @cascading_height = @current_depth
        @current_depth = nil
      end
    end

    def next_column_left
      @current_depth = @columns[@next_column_index].bottom
      # Go to leftmost column that has @current_depth as bottom
      while @columns[@next_column_index - 1].bottom == @current_depth
        @next_column_index -= 1
      end
    end

    # Parse a string representing a cave.
    def build_columns(src)
      lines = src.lines.to_a
      @map = lines
      lines = remove_borders(lines)
      column_heights = build_column_heights(lines)
      cols = column_heights.map { |ch| Column.new(ch) }
      cols
    end

    def build_column_heights(lines)
      heights = [0] * lines[0].length
      lines.each do |line|
        line.chars.with_index do |c, i|
          heights[i] += 1 if c != '#'
        end
      end
      heights
    end

    def remove_borders(lines)
      lines.delete_at(0)
      lines.delete_at(-1)
      lines.each { |l| l.chomp! "#\n" }
      lines
    end
  end
end