int_test.cr

require "benchmark"

lib IntIntrinsics
  {% for oper in %i(sadd ssub smul) %}
    {% for name, type in {i16: Int16, i32: Int32, i64: Int64} %}
      fun {{oper.id}}_with_overflow_{{name}} = "llvm.{{oper.id}}.with.overflow.{{name}}"(a : {{type}}, b : {{type}}) : { {{type}}, Bool }
    {% end %}
  {% end %}
  {% for oper in %i(uadd usub umul) %}
    {% for name, type in {i16: UInt16, i32: UInt32, i64: UInt64} %}
      fun {{oper.id}}_with_overflow_{{name}} = "llvm.{{oper.id}}.with.overflow.{{name}}"(a : {{type}}, b : {{type}}) : { {{type}}, Bool }
    {% end %}
  {% end %}
end

class IntegerOverflow < Exception
  def initialize(message = "Integer overflow")
    super(message)
  end
end

{% for name, type in {i16: Int16, i32: Int32, i64: Int64} %}
  struct {{type}}
    {% for oper in %i(add sub mul) %}
      @[AlwaysInline]
      def {{oper.id}}_with_overflow(other : {{type}}) : {{type}}
        result, overflow = IntIntrinsics.s{{oper.id}}_with_overflow_{{name}}(self, other)
        return result unless overflow
        raise IntegerOverflow.new
      end
    {% end %}

    @[AlwaysInline]
    def div_with_overflow(other : {{type}}) : {{type}}
      raise IntegerOverflow.new if (self == -1 || other == -1) && (self == {{type}}::MIN || other == {{type}}::MIN)
      self / other
    end
  end
{% end %}

{% for name, type in {i16: UInt16, i32: UInt32, i64: UInt64} %}
  struct {{type}}
    {% for oper in %i(add sub mul) %}
      @[AlwaysInline]
      def {{oper.id}}_with_overflow(other : {{type}}) : {{type}}
        result, overflow = IntIntrinsics.u{{oper.id}}_with_overflow_{{name}}(self, other)
        return result unless overflow
        raise IntegerOverflow.new
      end
    {% end %}
  end
{% end %}

def fibonacci(n)
  return n if n < 2
  return fibonacci(n - 2) + fibonacci(n - 1)
end

def fibonacci_safe(n)
  return n if n < 2
  return fibonacci_safe(n.sub_with_overflow(2))
    .add_with_overflow(fibonacci_safe(n.sub_with_overflow(1)))
end

COUNT = 30

Benchmark.ips(warmup: 2, calculation: 5) do |x|
  x.report("Fibonacci unsafe") {
    t = 0
    COUNT.times { |n| t += fibonacci(n) }
    puts t if t == 0
  }

  x.report("Fibonacci safe") {
    t = 0
    COUNT.times { |n| t = t.add_with_overflow(fibonacci_safe(n)) }
    puts t if t == 0
  }
end

#######################################################


struct Matrix(T)
  include Enumerable(T)

  alias IntRange = Range(Int32, Int32)


  getter cols : Int32
  getter rows : Int32

  def initialize(@rows, @cols, block : (Int32, Int32 -> T))
    @array = Array(T).new(@rows * @cols) { |i| block.call(*index_to_coords(i)) }
  end

  def self.new(rows, cols = rows, &block : (Int32, Int32 -> T))
    new(rows, cols, block)
  end

  def self.identity(size)
    new(size, size) { |r, c| (r == c ? 1 : 0) }
  end

  def self.[](*rows : Enumerable(T))
    new(rows.size, rows[0].size) { |r, c| rows[r][c] }
  end

  @[AlwaysInline]
  def [](row : Int32, col : Int32) : T
    check_index_out_of_bounds(row, col)
    @array[coords_to_index(row, col)]
  end

  @[AlwaysInline]
  def []=(row : Int32, col : Int32, val : T)
    check_index_out_of_bounds(row, col)
    @array[coords_to_index(row, col)] = val
  end

  def column(col)
    check_index_out_of_bounds(0, col)
    Array.new(rows) { |r| self[r, col] }
  end

  def row(row)
    check_index_out_of_bounds(row, 0)
    Array.new(cols) { |c| self[row, c] }
  end

  @[AlwaysInline]
  def *(m : Matrix)
    Matrix.new(rows, m.cols) do |i, j|
      (0...cols).reduce(0) do |vij, k|
        vij + self[i, k] * m[k, j]
      end
    end
  end

  @[AlwaysInline]
  def mul_safe(m : Matrix)
    Matrix.new(rows, m.cols) do |i, j|
      (0...cols).reduce(0) do |vij, k|
        vij.add_with_overflow(self[i, k].mul_with_overflow(m[k, j]))
      end
    end
  end


  def **(exp : Int)
    result = Matrix.identity(@rows)
    m = self
    while exp > 0
      result *= m if exp & 0b1 == 0b1
      m *= m
      exp >>= 1
    end
    result
  end

  def exp_safe(exp : Int)
    result = Matrix.identity(@rows)
    m = self
    while exp > 0
      result = result.mul_safe(m) if exp & 0b1 == 0b1
      m = m.mul_safe(m)
      exp >>= 1
    end
    result
  end


  @[AlwaysInline]
  private def index_to_coords(idx : Int32)
    {idx / @cols, idx % @cols}
  end

  @[AlwaysInline]
  private def coords_to_index(row : Int32, col : Int32)
    row * @cols + col
  end

  @[AlwaysInline]
  private def check_index_out_of_bounds(row : Int32, col : Int32)
    raise IndexError.new if col >= cols || col < 0 || row >= rows || row < 0
  end

  def to_s(io : IO)
    rows.times do |r|
      if rows == 1 ; io << '('
      elsif r == 0 ; io << '⎛'
      elsif r == rows - 1; io << '⎝'
      else io << '⎜'
      end
      io << ' '
      cols.times do |c|
        io << self[r, c]
        io << '\t' unless c == cols - 1
      end
      io << ' '
      if rows == 1; io << ')'
      elsif r == 0; io << '⎞' << '\n'
      elsif r == rows - 1; io << '⎠' << '\n'
      else io << '⎥' << '\n'
      end
    end
  end

  def transpose
    Matrix.new(cols, rows) { |r, c| self[c, r] }
  end
end


m = Matrix[
  [-3,0,0,0,1,1]
  [0,7,0,1,1,1],
  [0,0,1,0,0,-2],
  [0,2,0,1,1,0],
  [4,0,0,1,-1,0],
  [0,-2,5,-4,0,-1],
]

Benchmark.ips(warmup: 3, calculation: 8) do |x|
  x.report("Matrix exp unsafe") {
    t = nil
    (3...8).each { |n| t = m ** n }
    puts t if t && t == t.transpose
  }

  x.report("Matrix exp safe") {
    t = nil
    (3...8).each { |n| t = m.exp_safe(n) }
    puts t if t && t == t.transpose
  }
end