The Collapse Monad

collapse_monad.hs

import Control.Monad (ap, forM_)
import qualified Data.Map as M

-- The Collapse Monad
-- ------------------
-- The Collapse Monad allows collapsing a value with labelled superpositions
-- into a flat list of superposition-free values. It is like the List Monad,
-- except that, instead of always doing a cartesian product, it will perform
-- pairwise merges of distant parts of your program that are "entangled"
-- under the same label. Examples:
-- - collapse (&0{1 2}, 3)          == [(1,3), (2,3)]
-- - collapse (&0{1 2}, &0{3 4})    == [(1,3), (2,4)]
-- - collapse (&0{1 2}, &1{3 4})    == [(1,3), (1,4), (2,3), (2,4)]
-- - collapse (&0{1 2}, &1{3 4}, 5) == [(1,3,5), (1,4,5), (2,3,5), (2,4,5)]
-- Note how the second line above doesn't return the full cartesian product of
-- {1 2} and {3 4}; instead, it combines elements pairwise, because both
-- superpositions have the same label. I'm posting this file as a template to
-- recall later. This algorithm is used to flatten HVM3's results, which can be
-- superposed, back to a list of λ-terms without superpositions.

-- A bit-string
data Bin
  = O Bin
  | I Bin
  | E
  deriving Show

-- A Collapse is a tree of superposed values
data Collapse a = Sup Int (Collapse a) (Collapse a) | Val a

-- The Collapse Monad
-- Note: could be optimized using an IntMap instead of a List
bind :: Collapse a -> (a -> Collapse b) -> Collapse b
bind a f = fork a (repeat (\x -> x)) where
  fork (Val v)     paths = pass (f v) (map (\x -> x E) paths)
  fork (Sup k x y) paths = Sup k (fork x (mut k putO paths)) (fork y (mut k putI paths))
  pass (Val v)     paths = Val v
  pass (Sup k x y) paths = case paths !! k of
    E   -> Sup k x y
    O p -> pass x (mut k (\_->p) paths)
    I p -> pass y (mut k (\_->p) paths)
  putO bs = \x -> bs (O x)
  putI bs = \x -> bs (I x)

-- Collapses a Term and flattens into a list 
flatten :: Collapse a -> [a]
flatten (Val x)     = [x]
flatten (Sup _ x y) = flatten x ++ flatten y

-- Mutates an element at given index in a list
mut :: Int -> (a -> a) -> [a] -> [a]
mut 0 f (x:xs) = f x : xs
mut n f (x:xs) = x : mut (n-1) f xs
mut _ _ []     = []

instance Functor Collapse where
  fmap f (Val v)     = Val (f v)
  fmap f (Sup k x y) = Sup k (fmap f x) (fmap f y)

instance Applicative Collapse where
  pure  = Val
  (<*>) = ap

instance Monad Collapse where
  return = pure
  (>>=)  = bind

-- Example
-- -------
-- Collapsing a simple DSL with Superpositions

-- A simple DSL with tuples, numbers and superpositions
data Term
  = TNum Int           -- number
  | TTup Term Term     -- tuple
  | TSup Int Term Term -- superposition

-- Shows a term with superpositions
showTerm :: Term -> String
showTerm (TNum n)     = show n
showTerm (TTup x y)   = "(" ++ showTerm x ++ "," ++ showTerm y ++ ")"
showTerm (TSup k a b) = "&" ++ show k ++ "{" ++ showTerm a ++ " " ++ showTerm b ++ "}"

-- A Collapser for our DSL
-- - On normal values, we just use `<-`
-- - On superpositons, we return a Sup
collapse :: Term -> Collapse Term
collapse (TNum x) = do
  return $ TNum x
collapse (TTup l r) = do
  l <- collapse l
  r <- collapse r
  return $ TTup l r
collapse (TSup k a b) =
  Sup k (collapse a) (collapse b)

-- Test program
main :: IO ()
main = do
  -- Builds the term ((&0{1 2},&1{3 4}),(&0{5 6},&1{7 8}))
  let tup0 = TTup (TSup 0 (TNum 1) (TNum 2)) (TSup 1 (TNum 3) (TNum 4))
  let tup1 = TTup (TSup 0 (TNum 5) (TNum 6)) (TSup 1 (TNum 7) (TNum 8))
  let term = TTup tup0 tup1
  -- Collapses it
  let coll = flatten (collapse term)
  -- Prints the sup-free term list
  forM_ coll $ \ term ->
    putStrLn $ showTerm term

Note: it seems like the code above had a small bug, here's the version we're currently using on HVM3:

module Collapse where

import Control.Monad (ap, forM, forM_)
import qualified Data.IntMap.Strict as IM

import Type

-- The Collapse Monad 
-- ------------------
-- See: https://gist.github.com/VictorTaelin/60d3bc72fb4edefecd42095e44138b41

-- A bit-string
data Bin
  = O Bin
  | I Bin
  | E
  deriving Show

-- A Collapse is a tree of superposed values
data Collapse a = CSup Integer (Collapse a) (Collapse a) | CVal a | CEra
  deriving Show

cbind :: Collapse a -> (a -> Collapse b) -> Collapse b
cbind k f = fork k IM.empty where
  -- fork :: Collapse a -> IntMap (Bin -> Bin) -> Collapse b
  fork CEra         paths = CEra
  fork (CVal v)     paths = pass (f v) (IM.map (\x -> x E) paths)
  fork (CSup k x y) paths =
    let lft = fork x $ IM.alter (\x -> Just (maybe (putO id) putO x)) (fromIntegral k) paths in
    let rgt = fork y $ IM.alter (\x -> Just (maybe (putI id) putI x)) (fromIntegral k) paths in
    CSup k lft rgt 
  -- pass :: Collapse b -> IntMap Bin -> Collapse b
  pass CEra         paths = CEra
  pass (CVal v)     paths = CVal v
  pass (CSup k x y) paths = case IM.lookup (fromIntegral k) paths of
    Just (O p) -> pass x (IM.insert (fromIntegral k) p paths)
    Just (I p) -> pass y (IM.insert (fromIntegral k) p paths)
    Just E     -> CSup k (pass x paths) (pass y paths)
    Nothing    -> CSup k (pass x paths) (pass y paths)
  -- putO :: (Bin -> Bin) -> (Bin -> Bin)
  putO bs = \x -> bs (O x)
  -- putI :: (Bin -> Bin) -> (Bin -> Bin) 
  putI bs = \x -> bs (I x)

instance Functor Collapse where
  fmap f (CVal v)     = CVal (f v)
  fmap f (CSup k x y) = CSup k (fmap f x) (fmap f y)
  fmap _ CEra         = CEra

instance Applicative Collapse where
  pure  = CVal
  (<*>) = ap

instance Monad Collapse where
  return = pure
  (>>=)  = cbind

-- Simple Queue
-- ------------
-- Allows pushing to an end, and popping from another.
-- Simple purely functional implementation.
-- Includes sqPop and sqPut.

data SQ a = SQ [a] [a]

sqPop :: SQ a -> Maybe (a, SQ a)
sqPop (SQ [] [])     = Nothing
sqPop (SQ [] ys)     = sqPop (SQ (reverse ys) [])
sqPop (SQ (x:xs) ys) = Just (x, SQ xs ys)

sqPut :: a -> SQ a -> SQ a
sqPut x (SQ xs ys) = SQ xs (x:ys)

-- Flattener
-- ---------

flatten :: Collapse a -> [a]
flatten term = go term (SQ [] [] :: SQ (Collapse a)) where
  go (CSup k a b) sq = go CEra (sqPut b $ sqPut a $ sq)
  go (CVal x)     sq = x : go CEra sq
  go CEra         sq = case sqPop sq of
    Just (v,sq) -> go v sq
    Nothing     -> []