Sticking Functions Where They Donʼt Belong

Origin Story: Toy WASM Backend

I have been experimenting with making a silly little functional programming language that would compile to WASM, using GC types. (WASM GC provides structs and arrays with garbage collected references, integrated with the host runtime for references across the boundary: externref. Not that I wouldnʼt want to make my own GC sometimeee, but it is worth trying WASM GC out especially for prototyping.)

I really really want types in the language to have great freedom to choose how they are encoded, in types and in code. For example, the surface language provides basic syntax for arrays, text, and ADTs, but we should be able to encode enums as i8 or i31i31 shaves off a bit from i32 so it is colocatable with GC pointers or something. And maybe we want to have bitarrays encoded intelligently as bytearrays (array i8).

I donʼt want to maintain a registry of the custom types I add. I donʼt want to infect codegen with knowledge of specialized types and compilation. The basic codegen is so simple, it should stay that way. Keep the essence pure.

But I also donʼt want to give up on my expressions being pure data, an AST that I can inspect and compare and feed to property tests or something.

What am I to do?

The Existential Pattern

Letʼs smuggle functions into otherwise pure data, with a form of defunctionalization!

The magic of typeclasses is that they provide global, coherent dictionaries, keyed to types. Because they are global and coherent, they donʼt contain data themselves: all that you need to know is what type they have. Theoretically.

And the best part? Typeclass dictionaries contains functions, for free.

NFData isnʼt a problem for typeclass dictionaries: there are no unexpected thunks or bottoms waiting in global typeclass dictionaries. Eq and even Ord are already provided on TypeRep, the runtime representation of Typeable types (which is automatically derived). And Show isnʼt a problem (but Read will be).

First we make a typeclass synonym that bundles the nice things we want.

-- Nice properties for existentially hidden data
type Existentiable hidden = (Typeable hidden, NFData hidden, Ord hidden, Show hidden)

Next we make a typeclass that contains the functions we want to smuggle through. How about this one for injecting arbitrary foreign semantics into code.

class Existentiable sem => Semantics sem where
  semTyp :: sem -> STyp
  semCode :: sem -> WASM
  semEval :: sem -> [Expr] -> Map Name Expr -> Maybe Expr
  semExprs :: sem -> [Expr]
  semEffects :: sem -> Effects

The semantics type can be any data type that captures a behavior we want to see. Then once we have that data, we can project out the behaviors we care about: the type it has, some way to generate code, maybe simulate its evaluation, and an accounting of the effects it has, for static analysis purposes.

These are basically hooks into our compiler, places where we can extensibly inject arbitrary new behaviors.

Then we make an existential type that wraps any such data with its corresponding typeclass dictionary, including Typeable, Eq, and so on, up to the actual Semantics we want.

data ForeignSemantics = forall sem. Semantics sem => ForeignSemantics sem

Note that all of the typeclass methods are projective (they receive a sem and do not produce it), so we can implement Semantics for this abstract ForeignSemantics existential type too:

instance Semantics ForeignSemantics where
  semTyp (ForeignSemantics sem) = semTyp sem
  semCode (ForeignSemantics sem) = semCode sem
  semEval (ForeignSemantics sem) = semEval sem
  semExprs (ForeignSemantics sem) = semExprs sem
  semEffects (ForeignSemantics sem) = semEffects sem

To do this, we have to implement the Existentiable constituents.

Show is straightforward, just unwrap it and apply show. NFData also just forwards to the wrapped NFData – the dictionary, although it contains functions, does not need to be normalized for the reasons I discussed earlier: it is just a global typeclass dictionary.

Eq and Ord require comparing the runtime type reflection that is bundled in Typeable sem between the two values it gets. And of course Typeable ForeignSemantics comes for free – magic fairy dust courtesy of GHC.

instance Show ForeignSemantics where
  show (ForeignSemantics spec) = show spec
instance NFData ForeignSemantics where
  rnf (ForeignSemantics spec) = rnf spec

instance Eq ForeignSemantics where
  ForeignSemantics spec0 == ForeignSemantics spec2
    | Just spec1 <- Typeable.cast spec0 = spec1 == spec2
    | otherwise = False
instance Ord ForeignSemantics where
  ForeignSemantics spec0 `compare` ForeignSemantics spec2
    | Just spec1 <- Typeable.cast spec0 = spec1 `compare` spec2
    | otherwise = Typeable.typeOf spec0 `compare` Typeable.typeOf spec2

Now we get to use ForeignSemantics in the main AST for expressions, or wherever else it is useful, without spoiling deriving for Show or NFData or Eq or Ord.

-- The AST for expressions
data Expr
  -- Literals
  = EF64 !Double
  | EF32 !Float
  | EU64 !Word64
  | EU32 !Word32
  | ES64 !Int64
  | ES32 !Int32
  | ETxt !Text

  -- A let expression: multiple names to bind in order (not recursive let)
  | ELet [(Expr, Bind)] Expr
  -- A variable reference
  | EVar STyp !Name
  -- An n-ary lambda expression
  | EFun STyp [Bind] Expr
  -- Function application (n-ary)
  | EApp Expr [Expr]

  .
  .
  .

  -- Foreign code
  | EForeign ForeignSemantics
  deriving stock (Generic, Show, Eq, Ord)
  deriving anyclass (NFData)

Finally we can provide some helpful pattern synonyms. Since we bundled Typeable with this existential type, we can compare it against a specific type and see if it was that type, using Typeable.cast. (This requires a view pattern, which I am not really fond of, but is incredibly convenient in this situation.) The pattern synonym now works for destructing concrete types.

For convenience, we might as well repeat it for the AST type too. Call it ESem, since every Expr constructor starts with E.

pattern Sem :: Semantics sem => Semantics sem => sem -> ForeignSemantics
pattern Sem sem <- ForeignSemantics (Typeable.cast -> Just (sem :: sem)) where
  Sem sem = ForeignSemantics sem

pattern ESem :: Semantics sem => Semantics sem => sem -> Expr
pattern ESem sem = EForeign (Sem sem)

This is nice so that the library itself can provide instances and it still gets to match on them concretely when it needs to. For instance, STyp is another existential for synthetic types.

Now we can create instances and they just magically work. A minimal data type for a foreign call consists of its type, side effects that optimizations may need to know about, and the name of the WASM function or import to call.

data ForeignCall = ForeignCall STyp Effects Name
  deriving stock (Eq, Ord, Show, Generic)
  deriving anyclass (NFData)

instance Semantics ForeignCall where
  semTyp (ForeignCall t _ _) = t
  semCode (ForeignCall _ _ name) = tell $ SEXP "call" [ wasmID name ]
  semEval _ _ _ = Nothing
  semExprs _ = []
  semEffects (ForeignCall _ ef _) = ef

Et voilà, it is automatically hooked into the compiler and we can summon any function up with the magic expression ESem (ForeignCall (itsType) mempty "functionToCall").

Conclusion

Defunctionalization is usually either a conscious, excruciatingly explicit design choice, or something that is performed externally on a closed codebase, where all the functions are known.

This form of defunctionalization is extensible: you can write a Haskell library that contains some instances, and users can bring their own instances to the table, without any coordination.

You get to add extensible behaviors to any part of your code or data, without compromising much on what it means to be data, or tracking a registry of extensions, or trying to predict what your users (or you in the future) might want to do with it.

It requires some effort to determine what data is essential to the functions you want for an instance, but hopefully it is not too onerous and provides some clarity. It is better than the alternatives.