IO monad in Lua

monty-hall.lua

#!/usr/bin/env lua

local IO

IO = {
  -- Class method. Create an IO that does nothing and returns x.
  unit = function(x)
    return IO(function() return x end)
  end,

  -- Class method. Return a function which is like f, except it returns an IO
  -- that represents f's action.
  wrap = function(f)
    return function(...)
      local args <const> = table.pack(...)
      return IO(function() return f(table.unpack(args, 1, args.n)) end)
    end
  end,

  bind = function(self, f)
    return IO(function() return f(self())() end)
  end,

  -- Bind two IOs together, discarding a's value. Equivalent to
  -- a:bind(const(b)).
  seq = function(a, b)
    return IO(
      function()
	a()
	return b()
      end)
  end,

  rep = function(self, runs, f, init)
    local function loop(runs, x)
      if runs < 1 then
	return x
      end
      return loop(runs - 1, f(x, self()))
    end
    return IO(function() return loop(runs, init) end)
  end,
}

setmetatable(
  IO,
  {
    -- Construct a primitive IO with an impure function that takes no arguments.
    -- Equivalant to Haskell's 'ccall'.
    __call = function(cls, effect)
      local o <const> = {}
      setmetatable(o, {__index = cls, __call = effect})
      return o
    end,
  })

local random <const> = IO.wrap(math.random)
local print <const> = IO.wrap(print)

-- Based on https://www.youtube.com/watch?v=r6qg4NaUvfI&t=90s

local function simulateMonty(doorCount, swap)
  return
    random(doorCount)
    :bind(
      function(carDoor)
	return IO.unit(not not swap ~= (carDoor == 1))
      end)
end

-- Not a side-effect. It starts the program off at a known state.
math.randomseed(0)

local runs <const> = 1000000
local doors <const> = 26

local function simulate(swap)
  return simulateMonty(doors, swap):rep(
    runs,
    function(wins, win)
      if win then
	return wins + 1
      end
      return wins
    end,
    0)
end

-- Main program

simulate(false)
:bind(
  function(wins)
    return print(
      ('We won %d times out of %d without swapping.'):format(wins, runs))
  end)
:seq(
  simulate(true)
  :bind(
    function(wins)
      return print(
	('We won %d times out of %d while swapping.'):format(wins, runs))
    end))
()