Arrowized EDSL for Writing Lower Level Code

Question

Most of the motivating examples that are normally provided for Arrows are showing how more complex computational systems can be built on top of Hask (e.g. Kleisli categories for effects, Arrowized FRP etc.) Has there been any work done on using Arrows to write lower level code (e.g. Assembly, Javascript)? While this might not fit perfectly with the standard definition of Arrow (esp. arr :: (a -> b) -> cat a b), it seems that Arrows form a strong basis for some sort of concatenative programming.

Answer 1

Setting the bar

We are entering a programming limbo competition to see how low we can go with arrows. Against the blare from the audience, the judges are asking for our starting bar height. Lofted by the fans, we select the crowd favorite low level virtual machine as our target height. For the technical portion of the performance, we will implement a compiler for the primitive recursive functions defined via an ArrowLike interface I previously described.

module Control.PrimRec (
    ArrowLike (..),
    PrimRec (..),
    module Control.Category,
    module Data.Nat
) where

import Control.Category
import Data.Nat

import Prelude hiding (id, (.), fst, snd, succ)
import qualified Prelude (fst, snd)

class Category a => ArrowLike a where
    fst   :: a (b, d) b
    snd   :: a (d, b) b
    (&&&) :: a b c -> a b c' -> a b (c,c')

    first :: a b c -> a (b, d) (c, d)
    first = (*** id)

    second :: a b c -> a (d,b) (d,c)
    second = (id ***)

    (***) :: a b c -> a b' c' -> a (b,b') (c,c')
    f *** g = (f . fst) &&& (g . snd)

class ArrowLike a => PrimRec a where
    zero :: a b   Nat
    succ :: a Nat Nat
    prec :: a e c -> a (c, (Nat,e)) c -> a (Nat, e) c

Our goal is to make a Category that allows us to compose LLVM instructions together. We will also be providing an ArrowLike interface to juggle registers and the PrimRec interface to define functions on natural numbers.

Equipment check

The referee asks to see the equipment we'll be bringing onto the floor. We have two main challenges we're going to face. Even without arr introducing arbitrary functions, the types that our LLVM compiler can operate on will be quite constrained compared to all of Hask. The second challenge will be getting our LLVM instructions out of the program. Neither has much to do with Categorys, Arrows, or making a compiler, but our equipment bag and our code will be full of them to the point it will sometimes be hard to see the pieces we care about.

In our equipment bag, we have packed a host of tools for working with types.

{-# LANGUAGE GADTs #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators #-}

import GHC.Exts (Constraint)
import Data.Proxy

The referee takes our requisition for ordinary parts.

import Data.Word
import Data.Char (ord)

import Control.PrimRec

import Prelude hiding (
    id, (.), fst, snd, succ,
    sequence, sequence_, foldr,
    add)

The llvm-general-pure package has an AST for LLVM. We can pretty print the AST for use with the llvm tools with either llvm-general or llvm-pp.

import LLVM.General.AST hiding (type')
import LLVM.General.AST.Global
import LLVM.General.AST.Type
import qualified LLVM.General.AST.Constant as C
import qualified LLVM.General.AST.IntegerPredicate as ICmp
import qualified LLVM.General.AST.CallingConvention as CallingConvention

import LLVM.General.Pretty

We accept the standard issue of Applicative and Monad tools

import Data.Monoid
import Data.Foldable
import Data.Traversable

import Control.Applicative
import Control.Monad (forever)
import Control.Monad.Trans.Class
import Control.Monad.Trans.Writer.Strict (tell)
import Data.Functor.Identity

and bring a hoist in on a dolly.

import Control.Monad.Morph

A bunch of pipes clang around in the bottom of our bag.

import Pipes hiding (Proxy, void)
import qualified Pipes as P
import qualified Pipes.Prelude as P
import Pipes.Lift (runWriterP)

After inspecting the pipes to make sure none of them can be confused with the regulation limbo bar, the referee clears us to start coding.

Compiler Architecture

With a flourish and a clang, we define the state of a compiler.

type Build w = Pipe Name w

After puzzled gasps from the crowd, we define the two main operations that happen in a compiler: allocating a new symbol and emitting an instruction.

getName :: (Monad m) => Build w m (Name)
getName = await

instr :: (Monad m) => Named Instruction -> Build (Named Instruction) m ()
instr = yield

When our compiler is building something other than Named Instructions we can use yield to emit those instead.

Our compiler pieces are arrows that point from operands (registers or constants) to side effects in the compiler and new registers or constants. They are the Kleisli arrows of the compiled program: Operands x -> Build (Named Instruction) m (Operands y). We identify the compiler side effects of making new symbols and emitting instructions with side effects in the compiled program. The compiled program will also have global definitions, so we wrap the Kleisli arrows in the category of emitting global definitions. Unlike most arrows, due to the low level we are aiming at we can't accommodate arbitrary types. The arrows from registers to registers + side effects can only work on Registerable types that can be stored in LLVM registers or otherwise referred to by registers.

data RegisterArrow m x y where
    RegisterArrow :: (Registerable x, Registerable y) =>
                     (
                        Build Definition m (
                            Operands x ->
                            Build (Named Instruction) m (Operands y)
                        )
                     ) -> RegisterArrow m x y

Which types are Registerable are described in the next section.

Going up to go down

On our first try, we jump over the limbo bar. We want to write low level code, but, because the low level types don't match perfectly the Haskell types of Category, we need to go up first. We will need to juggle types and constraints to stay compatible with Category. It would be easier to ditch compatibility with Category in general and use leftaroundabout's constrained-categories. However, when we stay with Category we will smash your desire to "abstract away allowing for a variety of backends" - we will have the backend of ordinary Haskell functions -> described earlier, the backend of primitively recursive Haskell functions previously described, and our new LLVM compiler backend. So, first we go up and talk about constraints and kinds.

The Registerable types are those that can be associated with a representation in registers RegisterRep that is Traversable and supports application (Apply is Applicative without pure). They also must have an LLVM Type associated with each register.

class (Traversable (RegisterRep a), Apply (RegisterRep a)) => Registerable a where
    type RegisterRep a :: * -> *
    type RegisterableCtx a :: Constraint
    registerableDict :: Proxy a -> RegisterableDict a 
    types :: Proxy a -> Registers a Type

RegisterableCtx and RegisterableDict are used for inductive proofs later.

The unit type takes no registers. Its representation is a Proxy which holds nothing.

instance Registerable () where
    type RegisterRep () = Proxy
    type RegisterableCtx () = ()
    registerableDict _ = Dict
    types _ = Registers Proxy

A natural number can be stored in a register. Its representation is an Identity that stores one thing, and the type of that one thing is a 64 bit integer.

instance Registerable Nat where
    type RegisterRep Nat = Identity
    type RegisterableCtx Nat = ()
    registerableDict _ = Dict
    types _ = Registers . Identity $ IntegerType 64

A tuple can be stored in the registers for the two things it stores. Its representation is the product :*: of the representations of the two things it stores. Its types are the product of the type of the two things it stores. The tuple also introduces RegisterableCtx - in order to do everything with a tuple we need to know that both sides of it are Registerable.

instance (Registerable a, Registerable b) => Registerable (a, b) where
    type RegisterRep (a, b) = Registers a :*: Registers b
    type RegisterableCtx (a, b) = (Registerable a, Registerable b)
    registerableDict _ = Dict
    types _ = Registers $ types (Proxy :: Proxy a) :*: types (Proxy :: Proxy b)

We can define a Functor that has the shape of the registers for a Registerable type.

data Registers r a where
    Registers :: Registerable r => RegisterRep r a -> Registers r a

Because each Registerable type's representation is Traversable and has an Apply, we can define instances for the same for Registers.

instance Functor (Registers r) where
    fmap f (Registers xs) = Registers (fmap f xs)

instance Foldable (Registers r) where
    foldr f z (Registers xs) = foldr f z xs

instance Traversable (Registers r) where
    traverse f (Registers xs) = fmap Registers (traverse f xs)

instance Apply (Registers r) where
    Registers f <.> Registers x = Registers (f <.> x)

The Operands for a type we used earlier are just a structure with the same shape as the registers that hold that type, but holding an Operand in each place.

 type Operands f = Registers f Operand

Since the shape of registers can be traversed, they can be numbered in order.

number :: (Enum e, Traversable t) => (a -> e -> b) -> t a -> t b
number f = snd . mapAccumL (\(h:t) a -> (t, f a h)) [toEnum 0..]

Another lap...

With the fun part of the high level type programming out of the way, keeping track of what is Registerable becomes a slog. Bear in mind that none of this is visible to a user of the library, they only see the Category, ArrowLike and PrimRec type classes.

A RegisterableDict holds both the proof that a type is Registerable and whatever RegisterableCtx proof the type needs.

type RegisterableDict a = Dict (Registerable a, RegisterableCtx a)

A Dict holds the dictionary for a constraint. The dictionary is introduced to ghc when we pattern match on Dict. In order to construct a Dict, the constraint needs to be in scope in ghc. Registers and RegisterArrow earlier are also carrying around dictionaries, which are a boon when deconstructing and an obstacle when constructing.

data Dict c where
    Dict :: c => Dict c

We can now define something equivalent to RegisterArrow from the first section that can be defined for all types. Instead of constraining the types, we demand proof that the constraints are met for the input type, in the form of a RegisterableDict, before handing out the RegisterArrow that is stored inside. We will inductively prove from the input that the constraints are also met everywhere else.

data PRFCompiled m a b where
    BlockLike :: (RegisterableDict a -> RegisterArrow m a b) -> PRFCompiled m a b

To help keep track of the dictionaries, we'll make a few tools. rarrowDict recovers all of the constraints known about directly from a RegisterArrow

rarrowDict :: forall m x y. RegisterArrow m x y -> Dict (Registerable x, Registerable y, RegisterableCtx x, RegisterableCtx y)
rarrowDict (RegisterArrow _) =
    case registerableDict (Proxy :: Proxy x)
    of Dict ->
        case registerableDict (Proxy :: Proxy y)
        of Dict -> Dict

fstDict and sndDict demonstrate that if a tuple was Registerable then both of its components were.

fstDict :: forall a b. RegisterableDict (a, b) -> RegisterableDict a
fstDict Dict = case registerableDict (Proxy :: Proxy a) of Dict -> Dict

sndDict :: forall a b. RegisterableDict (a, b) -> RegisterableDict b
sndDict Dict = case registerableDict (Proxy :: Proxy b) of Dict -> Dict

Over and under at the same time

The compiler itself goes over the types and under the bar at the same time. It produces the low level instructions and builds the inductive proofs that each type is Registerable.

Each of the following instances is read most easily from inside to outside.

Kleisli arrows on registers form a Category. The identity id is to return the register the input was in. It produces no definitions, so in the category of definitions, it can just be returned. If passed in a Dict for the Registerable input, we know the output (which is the same) is also Registerable and can therefore build a RegisterArrow.

Composition of Kleisli arrows on registers, \a -> g a >>= f, says do all the side effects of the first one, pass the registers the result is in to the second one, do all of the side effects of the second one, and return the registers its result is in. Each component could also be producing definitions in the category of definitions, so we emit the definitions from both in order g <- mg; f <- mf; return .... Everything above the last three lines brings constraints into scope to inductively prove that constraints hold.

instance (Monad m) => Category (PRFCompiled m) where
    id  = BlockLike $ \Dict -> RegisterArrow . return $ return
    BlockLike df . BlockLike dg = BlockLike $ \Dict ->
        case dg Dict
        of rg@(RegisterArrow mg) ->
            case rarrowDict rg
            of Dict ->
                case df Dict
                of RegisterArrow mf -> RegisterArrow $ do
                    g <- mg
                    f <- mf
                    return (\a -> g a >>= f)

With just a category instance, we've written most of the compiler. We can put two computations together to build a new computation. Categorys form a strong basis for concatenative programming.

The ArrowLike instance does nothing but juggle which registers the compiler is referring to. That's also all Arrows do in Haskell - juggle which part of a structure you are working on with tuples. fst focuses on part of a structure of registers, snd focuses on the other part. &&& does two computations on the same set of registers and remembers the results of both.

instance (Monad m) => ArrowLike (PRFCompiled m) where
    fst = BlockLike $ \Dict -> RegisterArrow . return $ \(Registers (regs :*: _)) -> return regs
    snd = BlockLike $ \Dict -> RegisterArrow . return $ \(Registers (_ :*: regs)) -> return regs
    BlockLike df &&& BlockLike dg = BlockLike $ \Dict ->
        case (df Dict, dg Dict)
        of (RegisterArrow mf, RegisterArrow mg) -> RegisterArrow $ do
            f <- mf
            g <- mg
            return $ \regs -> do
                rf <- f regs
                rg <- g regs
                return $ Registers (rf :*: rg)

With Category and ArrowLike instances, we've written two thirds of our compiler and haven't even emitted a single instruction. All either instance has done is manipulated the state of the compiler or combined computations. Neither has performed any instructions. All of our computations come from the PrimRec instance, which introduces constructing and deconstructing natural numbers.

We construct natural numbers by constructing zero (bind an operand to the 0 constant) or compute the successor of a number (bind an operand to the result of adding 1 to the input operand).

instance (Monad m) => PrimRec (PRFCompiled m) where
    zero = BlockLike $ \Dict -> RegisterArrow . return $ \_ -> return . Registers . Identity . constant $ C.Int 64 0
    succ = BlockLike $ \Dict -> RegisterArrow . return $ regSucc
        where
            regSucc (Registers op) = (>>= return) . traverse opSucc $ Registers op
            opSucc op = bind i64 $ add op (constant $ C.Int 64 1)

We deconstruct natural numbers by primitive recursion, which we will implement naïvely and inefficiently in terms of recursion.

prec (BlockLike df) (BlockLike dg) = BlockLike $ \d@Dict ->
    case df $ sndDict d
    of (RegisterArrow mf) ->
        case dg Dict
        of (RegisterArrow mg) -> RegisterArrow $ do
            f <- mf
            g <- mg
            defineRecursive $ \go read ret -> do
                headName <- getName
                brName <- getName
                zeroName <- getName
                succName <- getName
                rs@(Registers (Registers (Identity n) :*: e)) <- block headName $ do
                    rs <- read
                    return (br brName,rs)
                block' brName $ do
                    cmp <- bind i1 $ icmp ICmp.EQ n (constant $ C.Int 64 0)
                    return (condbr cmp zeroName succName)
                block' zeroName $ do
                    c <- f e
                    ret c
                block' succName $ do
                    pred <- bind i64 $ sub n (constant $ C.Int 64 1)
                    c <- go (Registers (Registers (Identity pred) :*: e))
                    c' <- g (Registers (c :*: rs))
                    ret c'

We've just written a compiler for low level code directly inside its Category and ArrowLike instances. Silence fills the auditorium as the stunned band leader misses the que. Audience members in the front row faint.

Packing up

We start to casually pack our things, defining an extremely simple recursion to use in the definition of prec. Our "calling convention" per say will be to pass two pointers to a function, one pointing at memory from which it can read its arguments, and another pointing at memory to which it should write its results. Hackish and unsophisticated, the audience starts chatting through the denouement, but the judges are still eager to see if it all actually works.

defineRecursive :: forall x y m. (Registerable x, Registerable y, Monad m) =>
                    (
                        (Operands x -> Build (Named Instruction) m (Operands y))       ->  -- recursive call
                                       Build (Named Instruction) m (Operands x)        ->  -- read parameters 
                        (Operands y -> Build (Named Instruction) m (Named Terminator)) ->  -- return results
                        Build (BasicBlock) m ()                                            -- function body
                    ) ->
                    Build Definition m (
                        Operands x -> Build (Named Instruction) m (Operands y))            -- call function
defineRecursive def = do
    functionName <- getName
    inPtrName <- getName
    outPtrName <- getName
    let
        inType  = StructureType False . toList $ types (Proxy :: Proxy x)
        outType = StructureType False . toList $ types (Proxy :: Proxy y)
        outPtrType = ptr outType
        inPtrType  = ptr inType
        go  regs = do
            inPtr  <- bind (ptr inType)  $ alloca inType
            outPtr <- bind (ptr outType) $ alloca outType
            writePtr inPtr regs
            instr $ call
                        (constant $ C.GlobalReference (FunctionType void [ptr outType, ptr inType] False) functionName)
                        [outPtr, inPtr]
            readPtr outPtr
        ret regs = do
            writePtr (LocalReference outPtrType outPtrName) regs
            return (retVoid)
        read = readPtr (LocalReference inPtrType inPtrName)
    (blocks, _) <- collect (def go read ret)
    yield $ global $ define void functionName [(outPtrType, outPtrName), (inPtrType, inPtrName)] blocks
    return go

A heckler yells something about a stack frame per successor, "... 5 digits at most!".

Storing or retrieving every field of a Traversable structure in memory one after another is as straightforward as traversing it. We pack all our data into stack memory, not worrying about sharing or duplication. After all we were long since "liberate[d] from petty concerns about, e.g., the efficiency of hardware-based integers".

elemPtrs :: (Monad m, Traversable f) => Operand -> f Type -> Build (Named Instruction)  m (f Operand)
elemPtrs struct ts = do
    sequence $ number getElemPtr ts
    where
        getElemPtr t n = bind (ptr t) $ getelementptr struct [C.Int 32 0, C.Int 32 n]

readPtr :: forall r m. (Registerable r, Monad m) => Operand ->  Build (Named Instruction) m (Operands r)
readPtr struct = do
    let ts  = types (Proxy :: Proxy r)
    elems <- elemPtrs struct ts
    sequence $ (bind <$> ts) <.> (load <$> elems)

writePtr :: forall r m. (Registerable r, Monad m) => Operand -> Operands r -> Build (Named Instruction) m ()
writePtr struct ops = do
    let ts  = types (Proxy :: Proxy r)
    elems <- elemPtrs struct ts
    sequence_ $ instr . Do <$> (store <$> ops <.> elems)

As we put the last 64 bit integer into its stack allocated cell, the last of the audience trickles out of the stands.

Going sideways while going forward

A small crowd gathers around the hoist. "You never used this", one says. "We did, it just might not have made the jumbo-tron. We used it when we went deeper into a block".

Every time we yield or await the pipe, data moves sideways across the flow of computation. It comes in on one side when we await a new name and goes out the other side when we yield a result. Every time we bound an operand to an instruction we were doing both, but in the result of the monad only had to worry about the operands holding the result of the computation.

bind :: (Monad m) => Type -> Instruction -> Build (Named Instruction) m (Operand)
bind t instruction = do
    name <- getName
    instr $ name := instruction
    return (LocalReference t name)

To make a block we collect the yielded results of a sub computation.

block :: (Monad m) => Name -> Build (Named Instruction) m (Named Terminator, r) -> Build BasicBlock m r
block name definition = do
    (instructions, (terminator, r)) <- collect definition
    yield $ BasicBlock name instructions terminator
    return r

block' name = block name . (>>= \x -> return (x,()))

To collect the results of the sub computations, we take advantage of the fact that pipes are so pure that you can swap out the underlying monad provided you have a natural transformation forall x. m x -> n x. This is what hoist does; it lets us lift all of the underlying monad operations under a pipe to use the pipe over another transformer, in this case WriterT.

collect :: (Monad m) => Pipe a b m r -> Pipe a c m ([b], r)
collect subDef = do
    (r, w) <- runWriterP $
        hoist lift subDef >->
        forever (await >>= \x -> lift $ tell (++[x]))
    return (w [], r)

How low did we go?

The head judge presents a selection from our chosen bar height, and asks us to compile isOdd. The example code is written only in terms of the three interfaces Category, ArrowLike, and PrimRec.

match :: PrimRec a => a b c -> a (Nat, b) c -> a (Nat, b) c
match fz fs = prec fz (fs . snd)

one :: PrimRec a => a b Nat
one = succ . zero

isZero :: PrimRec a => a Nat Nat
isZero = match one zero . (id &&& id)

isOdd :: PrimRec a => a Nat Nat
isOdd = prec zero (isZero . fst) . (id &&& id)

The panel smirks at the extremely inefficient implementation of isZero.

define void @n1({i64}* %n3, {i64, i64}* %n2){
n4:
  %n8 = getelementptr inbounds {i64, i64}* %n2, i32 0, i32 0
  %n9 = getelementptr inbounds {i64, i64}* %n2, i32 0, i32 1
  %n10 = load i64* %n8
  %n11 = load i64* %n9
  br label %n5
n5:
  %n12 = icmp eq i64 %n10, 0
  br i1 %n12, label %n6, label %n7
n6:
  %n13 = add i64 0, 1
  %n14 = getelementptr inbounds {i64}* %n3, i32 0, i32 0
  store i64 %n13, i64* %n14
  ret void
n7:
  %n15 = sub i64 %n10, 1
  %n16 = alloca {i64, i64}
  %n17 = alloca {i64}
  %n18 = getelementptr inbounds {i64, i64}* %n16, i32 0, i32 0
  %n19 = getelementptr inbounds {i64, i64}* %n16, i32 0, i32 1
  store i64 %n15, i64* %n18
  store i64 %n11, i64* %n19
  call void @n1({i64}* %n17, {i64, i64}* %n16)
  %n20 = getelementptr inbounds {i64}* %n17, i32 0, i32 0
  %n21 = load i64* %n20
  %n22 = getelementptr inbounds {i64}* %n3, i32 0, i32 0
  store i64 0, i64* %n22
  ret void
}

One of the junior judges says, "It's technically correct" which a colleague remarks is the "worst kind of correct". The referee sits down at the console and tries some inputs:

123456
0

54321
1

654321
[Stack Overflow]

The head judge congratulates us on the low level of the bar, but suggests we should try something "more technical" for full marks.

Memorabilia

As we walk out through the lobby, we notice that that our complete code and output was printed in a special late edition of the CoPointed Sheet, available at the stadium for $12.50 and at the newsstand outside for $1.25.