assembly.rb

assembly.rb

# frozen_string_literal: true

require 'minitest/autorun'

# This is follow-up to https://medium.com/@kmadej/interviewing-in-ruby-a-minimal-assembly-vm-3486d113c4e9.

# Sample program:
#
# MOV A 10  # Store 10 in register A
# MOV B A   # Copy value from register A to register B
# ADD A B   # Add register A and B and store the result in register C
# MOV A 100 # Store 100 in register A
# DEC C     # Decrement value of register C
# DEC A     # Decrement value of register A
# JNZ 5     # Jump to line #5 if the value in register C is non-zero
# MUL A B   # Multiply values stored in environment A and B and store them in register C

# Usage:
#
# $> irb
# 2.5.3 :001 > require './assembly.rb'
#  => true
# 2.5.3 :002 > program = <<~ASM
# 2.5.3 :003">   MOV A 10
# 2.5.3 :004">   MOV B A
# 2.5.3 :005">   ADD A B
# 2.5.3 :006">   MOV A 100
# 2.5.3 :007">   DEC C
# 2.5.3 :008">   DEC A
# 2.5.3 :009">   JNZ 5
# 2.5.3 :010">   MUL A B
# 2.5.3 :011"> ASM
#  => "MOV A 10\nMOV B A\nADD A B\nMOV A 100\nDEC C\nDEC A\nJNZ 5\nMUL A B\n"
# 2.5.3 :012 > Assembly.new.execute(program)
#  => #<Env::Environment:0x00.. @pointer=8, @registers=#<Env::Registers:0x00.. @a=80, @b=10, @c=800>>

ParseError = Class.new(StandardError)

class Assembly
  def execute(program, environment = Env::Environment.new)
    operations = parse_program(program)
    run(operations, environment)
  end

  private

  def parse_program(program)
    program.split("\n").map { |line| parse_line(line) }
  end

  def parse_line(line)
    operation, *params = line.split

    case operation
    when 'MOV' then Operations::Mov.new(*params)
    when 'ADD' then Operations::Add.new(*params)
    when 'MUL' then Operations::Mul.new(*params)
    when 'DEC' then Operations::Dec.new(*params)
    when 'JNZ' then Operations::Jnz.new(*params)
    else raise ::ParseError, "Operation #{operation} not supported."
    end
  end

  def run(operations, environment)
    next_operation = operations[environment.pointer]
    return environment unless next_operation
    environment = next_operation.execute(environment)
    run(operations, environment)
  end
end

module Env
  class Environment
    REGISTERS_NAMES = ['A', 'B', 'C']

    def initialize(pointer: 0, a: 0, b: 0, c: 0)
      @pointer = pointer
      @registers = Registers.new(a: a, b: b, c: c)
    end

    attr_accessor :pointer

    def set_register_value(name, value)
      @registers.set(name, value)
    end

    def get_register_value(name)
      @registers.get(name)
    end

    def increment_pointer
      @pointer += 1
    end
  end

  class Registers
    def initialize(a: 0, b: 0, c: 0)
      @a = a
      @b = b
      @c = c
    end

    def set(name, value)
      send("#{name}=", value)
    end

    def get(name)
      send(name)
    end

    private

    attr_accessor :a, :b, :c
  end
end

module Operations
  class Mov
    def initialize(dest, src)
      @dest = Params::RegisterReference.build(dest)
      @src = Params::NumericalValueOrRegisterReference.build(src)
    end

    def execute(environment)
      environment.set_register_value(@dest.register_name, @src.value(environment))
      environment.increment_pointer
      environment
    end
  end

  class Add
    def initialize(x, y)
      @x = Params::NumericalValueOrRegisterReference.build(x)
      @y = Params::NumericalValueOrRegisterReference.build(y)
    end

    def execute(environment)
      environment.set_register_value(:c, @x.value(environment) + @y.value(environment))
      environment.increment_pointer
      environment
    end
  end

  class Mul
    def initialize(x, y)
      @x = Params::NumericalValueOrRegisterReference.build(x)
      @y = Params::NumericalValueOrRegisterReference.build(y)
    end

    def execute(environment)
      environment.set_register_value(:c, @x.value(environment) * @y.value(environment))
      environment.increment_pointer
      environment
    end
  end

  class Dec
    def initialize(x)
      @x = Params::RegisterReference.build(x)
    end

    def execute(environment)
      environment.set_register_value(@x.register_name, @x.value(environment) - 1)
      environment.increment_pointer
      environment
    end
  end

  class Jnz
    def initialize(x)
      @x = Params::LineNumber.build(x)
    end

    def execute(environment)
      if environment.get_register_value(:c).nonzero?
        environment.pointer = @x.value - 1
      else
        environment.increment_pointer
      end
      environment
    end
  end

  module Params
    class NumericalValueOrRegisterReference
      def self.build(param)
        if RegisterReference.valid?(param)
          RegisterReference.build(param)
        elsif NumericalValue.valid?(param)
          NumericalValue.build(param)
        else
          raise ::ParseError, "Invalid param: #{param} (expected numerical value or register reference)."
        end
      end
    end

    class NumericalValue
      def self.build(param)
        raise ::ParseError, "Invalid param: #{param} (expected numerical value)." unless valid?(param)
        new(param.to_i)
      end

      def self.valid?(param)
        param =~ /\d+/
      end

      def initialize(param)
        @param = param
      end

      def value(_registers)
        @param
      end
    end

    class RegisterReference
      def self.build(param)
        raise ::ParseError, "Invalid param: #{param} (exprected register reference)." unless valid?(param)
        new(param.downcase.to_sym)
      end

      def self.valid?(param)
        ::Env::Environment::REGISTERS_NAMES.include?(param)
      end

      def initialize(param)
        @param = param
      end

      def value(environment)
        environment.get_register_value(@param)
      end

      def register_name
        @param
      end
    end

    class LineNumber
      def self.build(param)
        raise ::ParseError, "Invalid param: #{param} (exprected line number)." unless valid?(param)
        new(param.to_i)
      end

      def self.valid?(param)
        param =~ /\d+/
      end

      def initialize(param)
        @param = param
      end

      def value
        @param
      end
    end
  end
end

class AssemblyTest < Minitest::Test
  def setup
    @assembly = Assembly.new
  end

  def test_empty_program
    environment = @assembly.execute('')

    assert_equal(0, environment.get_register_value(:a))
    assert_equal(0, environment.get_register_value(:b))
    assert_equal(0, environment.get_register_value(:c))
  end

  def test_unsupported_operation
    assert_raises(ParseError) { @assembly.execute('XXX') }
  end

  def test_mov_a_1
    environment = @assembly.execute('MOV A 1')

    assert_equal(1, environment.get_register_value(:a))
    assert_equal(0, environment.get_register_value(:b))
    assert_equal(0, environment.get_register_value(:c))
  end

  def test_mov_b_a
    environment = @assembly.execute('MOV B A', Env::Environment.new(a: 1))

    assert_equal(1, environment.get_register_value(:a))
    assert_equal(1, environment.get_register_value(:b))
    assert_equal(0, environment.get_register_value(:c))
  end

  def test_mov_1_a
    assert_raises(ParseError) { @assembly.execute('MOV 1 A') }
  end

  def test_multiline_program
    environment = @assembly.execute("MOV A 1\nMOV B 1")

    assert_equal(1, environment.get_register_value(:a))
    assert_equal(1, environment.get_register_value(:b))
    assert_equal(0, environment.get_register_value(:c))
  end

  def test_add_a_b
    environment = @assembly.execute('ADD A B', Env::Environment.new(a: 1, b: 2))

    assert_equal(1, environment.get_register_value(:a))
    assert_equal(2, environment.get_register_value(:b))
    assert_equal(3, environment.get_register_value(:c))
  end

  def test_add_1_a
    environment = @assembly.execute('ADD 1 A', Env::Environment.new(a: 2))

    assert_equal(2, environment.get_register_value(:a))
    assert_equal(0, environment.get_register_value(:b))
    assert_equal(3, environment.get_register_value(:c))
  end

  def test_add_a_1
    environment = @assembly.execute('ADD A 1', Env::Environment.new(a: 2))

    assert_equal(2, environment.get_register_value(:a))
    assert_equal(0, environment.get_register_value(:b))
    assert_equal(3, environment.get_register_value(:c))
  end

  def test_mul_a_b
    environment = @assembly.execute('MUL A B', Env::Environment.new(a: 2, b: 3))

    assert_equal(2, environment.get_register_value(:a))
    assert_equal(3, environment.get_register_value(:b))
    assert_equal(6, environment.get_register_value(:c))
  end

  def test_mul_2_a
    environment = @assembly.execute('MUL 2 A', Env::Environment.new(a: 3))

    assert_equal(3, environment.get_register_value(:a))
    assert_equal(0, environment.get_register_value(:b))
    assert_equal(6, environment.get_register_value(:c))
  end

  def test_mul_a_2
    environment = @assembly.execute('MUL A 2', Env::Environment.new(a: 3))

    assert_equal(3, environment.get_register_value(:a))
    assert_equal(0, environment.get_register_value(:b))
    assert_equal(6, environment.get_register_value(:c))
  end

  def test_dec_a
    environment = @assembly.execute('DEC A', Env::Environment.new(a: 2))

    assert_equal(1, environment.get_register_value(:a))
    assert_equal(0, environment.get_register_value(:b))
    assert_equal(0, environment.get_register_value(:c))
  end

  def test_jnz_3_executed
    environment = @assembly.execute('JNZ 3', Env::Environment.new(c: 1))

    assert_equal(2, environment.pointer)
  end

  def test_jnz_3_skipped
    environment = @assembly.execute('JNZ 3')

    assert_equal(1, environment.pointer)
  end

  def test_whole_program
    program = <<~ASM
      MOV A 10
      MOV B A
      ADD A B
      MOV A 100
      DEC C
      DEC A
      JNZ 5
      MUL A B
    ASM

    environment = @assembly.execute(program)

    assert_equal(80, environment.get_register_value(:a))
    assert_equal(10, environment.get_register_value(:b))
    assert_equal(800, environment.get_register_value(:c))
  end
end