1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
# typed: strict

require "stringio"

require "sorbet-runtime"

AnyIO = T.type_alias { T.any(IO, StringIO) }
Memory = T.type_alias { T::Array[T.nilable(Integer)] }

class Mode < T::Enum
  enums do
    Position = new
    Immediate = new
    Relative = new
  end

  extend T::Sig

  sig {params(value: String).returns(Mode)}
  def self.from(value)
    case value
    when ?0 then Mode::Position
    when ?1 then Mode::Immediate
    when ?2 then Mode::Relative
    else fail "unexpected mode: #{value}"
    end
  end
end

OPCODES = T.let({
  1  => ->(x, _, _, a, b, c) { x[c]  = x[a] + x[b]            }, # add
  2  => ->(x, _, _, a, b, c) { x[c]  = x[a] * x[b]            }, # multiply
  3  => ->(x, i, _, a)       { x[a]  = i.gets.to_i            }, # input
  4  => ->(x, _, o, a)       {         o.puts(x[a])           }, # output
  5  => ->(x, _, _, a, b)    { x.pc  = x[b] if x[a].nonzero?  }, # jump-if-true
  6  => ->(x, _, _, a, b)    { x.pc  = x[b] if x[a].zero?     }, # jump-if-false
  7  => ->(x, _, _, a, b, c) { x[c]  = (x[a] < x[b])  ? 1 : 0 }, # less than
  8  => ->(x, _, _, a, b, c) { x[c]  = (x[a] == x[b]) ? 1 : 0 }, # equals
  9  => ->(x, _, _, a)       { x.rb += x[a]                   }, # adjust relative base
  99 => ->(*)                {         throw :halt            }, # halt
}, T::Hash[Integer, T.untyped])

class Parameter
  extend T::Sig

  sig {params(value: T.any(Parameter, Integer)).returns(Parameter)}
  def self.from(value)
    case value
    when Parameter then value
    when Integer then new(Mode::Position, value)
    else T.absurd(value)
    end
  end

  sig {returns(Mode)}
  attr_reader :mode

  sig {returns(Integer)}
  attr_reader :value

  sig {params(mode: Mode, value: Integer).void}
  def initialize(mode, value)
    @mode = T.let(mode, Mode)
    @value = T.let(value, Integer)
  end
end

class Computer
  extend T::Sig

  sig {params(input: String).returns(Computer)}
  def self.from(input)
    new(input.split(?,).map(&:to_i))
  end

  sig {returns(Integer)}
  attr_accessor :pc, :rb

  sig {params(program: Memory).void}
  def initialize(program)
    @memory = T.let(program.dup, Memory)
    @pc = T.let(0, Integer)
    @rb = T.let(0, Integer)
  end

  sig {params(input: T.untyped, output: AnyIO).returns(Memory)}
  def run(input=STDIN, output=STDOUT)
    each(input, output).inject(nil) {|_,i| i }
  end

  sig {
    params(
      input: T.untyped,
      output: AnyIO,
      blk: T.nilable(T.proc.params(m: Memory).returns(T.nilable(Integer)))
    ).returns(T::Enumerator[Memory])
  }
  def each(input, output, &blk)
    return enum_for(T.must(__method__), input, output) unless block_given?

    catch(:halt) do
      loop do
        instruction = @memory[pc].to_s.rjust(5, ?0)
        opcode = OPCODES.fetch(instruction[-2..-1].to_i)
        self.pc += 1

        n = opcode.arity - 3 # subtract the computer, input, and output params
        args = (0...n).map {|i|
          mode = Mode.from(T.must(instruction[2-i]))
          value = @memory[pc + i] || 0
          Parameter.new(mode, value)
        }
        self.pc += n

        opcode.call(self, input, output, *args)

        yield @memory
      end
    end

    # Satisfy sorbet:
    #   https://sorbet-ruby.slack.com/archives/CHN2L03NH/p1575648549254600
    Enumerator.new {}
  end

  sig {params(parameter: Parameter).returns(Integer)}
  def [](parameter)
    parameter = Parameter.from(parameter)
    mode = parameter.mode
    case mode
    when Mode::Position  then @memory[parameter.value] || 0
    when Mode::Immediate then parameter.value
    when Mode::Relative  then @memory[rb + parameter.value] || 0
    else T.absurd(mode)
    end
  end
  alias_method :fetch, :[]

  sig {params(parameter: Parameter, value: Integer).void}
  def []=(parameter, value)

    mode = parameter.mode
    case mode
    when Mode::Position  then @memory[parameter.value] = value
    when Mode::Immediate then raise "writes should never be in immediate mode"
    when Mode::Relative then @memory[rb + parameter.value] = value
    else T.absurd(mode)
    end
  end
end

if __FILE__ == $0
  Computer.from(T.cast(ARGF, IO).read || "").run
end