Data.Edit.Tutorial

This is a short (?) tutorial describing how you can use the Edit module to help you with writing dataflow analysis code for a compiler. The example is a bit artificial for the sake of relative conciseness -- if you have a better suggestion, or find any mistakes, please let me know on the Github issue tracker.

TL;DR

Get a fixed point from applying a sequence of transformations.

import Data.Edit (Edit, edits, polish, (>=>))

mkAwesome1 :: Foo -> Maybe Foo
...
mkAwesomeN :: Foo -> Maybe Foo

mkAwesomeAny :: Foo -> Edit Foo
mkAwesomeAny
  = foldr (\f acc -> acc >=> (f `edits`)) pure
    [mkAwesome1, ..., mkAwesomeN]

mkAsAwesomeAsPossible :: Foo -> Foo
mkAsAwesomeAsPossible = polish mkAwesomeAny

Transform a recursive data structure, keeping track of whether it was changed or not, and feed the result to some high-level dataflow analysis function.

import DataFlowLibrary
import PlatedLibrary
import Data.Edit (Edit, edits, toMaybe)

instance FancyPlate Foo where ...

mkAwesome :: Foo -> Maybe Foo
mkAwesome = ...

mkTotallyAwesome :: Foo -> Edit Foo
mkTotallyAwesome = transformM (mkAwesome `edits`)

dataFlowAnalysis = dataFlowLibFn (toMaybe . mkTotallyAwesome)

Setup

The examples here use the Uniplate and Containers libraries. If you want to follow along as we proceed, you will want to supply the package flag tutorial and maybe read the docs in your browser.

If you're testing inside a cabal sandbox, this can be done using

cabal configure --flags="tutorial"
cabal build
cabal haddock

If you're using stack, the same can be done using:

stack build --flag=edit:tutorial
stack haddock --flag=edit:tutorial --open edit

Tutorial

Let's define a toy language L with Ints and addition.

newtype Ident = Ident String
  deriving (Show, Eq)

data Expr
  = Val Int
  | Var Ident
  | Add Expr Expr
  deriving (Show, Eq)

Q. How would you implement constant folding for the Expr type?

1. Write the recursion by hand. While this is easy enough to do since Expr only has a few constructors, this isn't very practical when you have lots of constructors. The exact point where you recognize that this is a recursive descent into unmaintainability depends on your personal boilerplate threshold.
2. Use recursion schemes and get lost in the unfathomable type errors (I'm half-joking). While this is a reasonable approach, we're not going to follow this here.
3. Use a generics library. For simplicity, we'll be using Uniplate here. The particular functions that are relevant at the moment are rewrite and transform. Let's use rewrite.

{-# LANGUAGE DeriveDataTypeable #-}

import Data.Data
import Data.Generics.Uniplate.Data

newtype Ident = Ident String
  deriving (Show, Eq, Typeable, Data)

data Expr
  = Val Int
  | Var Ident
  | Add Expr Expr
  deriving (Show, Eq, Typeable, Data)

constFold :: Expr -> Expr
constFold e = rewrite go e
  where
    go (Add (Val i) (Val j)) = Just (Val (i + j))
    go _ = Nothing

Test that the implementation works as expected.

>>> two = Add (Val 1) (Val 1)
>>> four = Add (Val 2) (Val 2)
>>> constFold (Add two four)
Val 6
>>> constFold (Add (Var "x") two)
Add (Var "x") (Val 2)

Let's say we add assignment statements to the language and write a function to do constant propagation. First we add a substitute function.

import Data.Map (Map)
import qualified Data.Map as Map

newtype Ident = Ident String
  deriving (Eq, Ord, Show, Typeable, Data)

substitute :: Map Ident Int -> Expr -> Expr
substitute m e = rewrite go e
  where
    go (Var x) = Val <$> Map.lookup x m
    go _ = Nothing

Let's test this out.

>>> x = Var (Ident "x")
>>> quadrupleX = Add x (Add x (Add x x))
>>> m1 = Map.fromList [(Ident "x", 5)]
>>> substitute m1 quadrupleX
Add (Val 5) (Add (Val 5) (Add (Val 5) (Val 5)))

Finally add in statements and a constant propagation function.

infix 9 :=
data Stmt = Ident := Expr
    deriving (Show)

constProp :: Map Ident Int -> Stmt -> (Map Ident Int, Stmt)
constProp map_ (var := expr) = (f map_, var := expr')
  where
    expr' = substitute map_ expr
    f = case expr' of
      Val x -> Map.insert var x
      _ -> Map.delete var -- delete old entry if var is re-defined

>>> x = Var (Ident "x")
>>> m1 = Map.fromList [(Ident "x", 5)]
>>> constProp m1 (Ident "y" := Var (Ident "x"))
(fromList [(Ident "x",5),(Ident "y",5)],Ident "y":=Val 5)

Now let's say we want to write two passes -- one for constant folding, one for constant propagation, and then iterate until no more optimization can be done (yes, this isn't an optimal strategy, but then this tutorial would be even longer :O).

However, the constFold function, as it stands, doesn't save the "information" whether it changed something or not. Consequently, we won't be able to tell if we hit the fixed point or not unless we do an equality check (which could be expensive if the expression trees are big). Time to finally use the Edit monad!

We can use the edits function, which converts a function f: a -> Maybe a to a function f' : a -> Edit a.

import Data.Edit

-- We don't have to alter the core logic here, neat!
constFold' :: Expr -> Edit Expr
constFold' = transformM (go `edits`)
 where
  go (Add (Val i) (Val j)) = Just (Val (i + j))
  go _ = Nothing

constFoldPass :: [Stmt] -> Edit [Stmt]
constFoldPass ss = traverse (\(v := e) -> (v :=) <$> constFold' e) ss

We also need slightly different versions of substitute and constProp. Here we use the extract function; it has the signature Edit a -> a. It is fine to throw away the 'Clean'/'Dirty' information when we are updating the map, because we are only interested in changes to the Stmt and don't care if the Map gets changed or not.

substitute' :: Map Ident Int -> Expr -> Edit Expr
substitute' m e = transformM (go `edits`) e
 where
  go (Var x) = Val <$> Map.lookup x m
  go _ = Nothing

constProp' :: Map Ident Int -> Stmt -> (Map Ident Int, Edit Stmt)
constProp' map_ (var := expr) = (f map_, (var :=) <$> expr')
 where
  expr' = substitute' map_ expr
  f = case extract expr' of
    Val x -> Map.insert var x
    _ -> id

Let's add a top-level function similar to constFoldPass.

Note: If you're unfamiliar with unfoldr, you can think of it as the opposite of foldr. foldr takes a list and a starting value and collapses it to a single value; unfoldr takes a starting value (often called a seed) and generates a list out of it.

import Data.List (unfoldr)

constPropPass :: [Stmt] -> Edit [Stmt]
constPropPass ss = sequence $ unfoldr go (Map.empty, ss)
 where
  go (_, []) = Nothing
  go (m, x:xs) = let (m', ex) = constProp' m x in Just (ex, (m', xs))

Finally putting all the pieces together. We can use the polish function to find the fixed point, which (in this case) is a fancy way of saying that we keep iterating until we have a Clean (unchanged) value.

constFoldAndPropPass :: [Stmt] -> [Stmt]
constFoldAndPropPass = polish (constFoldPass >=> constPropPass)

We're not done yet though! We still need to check that this works :P.

>>> [w, x, y] = map Ident ["w", "x", "y"]
>>> s1 = w := Add (Val 1) (Val 2)
>>> s2 = x := Add (Var w) (Var w)
>>> s3 = y := Add (Var w) (Add (Val 1) (Var x))
>>> s4 = x := Add (Var y) (Var y)
>>> s5 = y := Add (Var w) (Var x)
>>> constFoldAndPropPass [s1, s2, s3, s4, s5]
[Ident "w" := Val 3,Ident "x" := Val 6,Ident "y" := Val 10,Ident "x" := Val 20,Ident "y" := Val 23]

Yup, it works! For fun, let's see the transformation process in action. We can do this using the iterations function.

>>> pprint = putStr . unlines . map (unlines . map show)
>>> pprint $ iterations (constFoldPass >=> constPropPass) [s1, s2, s3, s4, s5]

The output shows the full history, with the final result that we obtained earlier at the end.

Ident "w" := Add (Val 1) (Val 2)
Ident "x" := Add (Var (Ident "w")) (Var (Ident "w"))
Ident "y" := Add (Var (Ident "w")) (Add (Val 1) (Var (Ident "x")))
Ident "x" := Add (Var (Ident "y")) (Var (Ident "y"))
Ident "y" := Add (Var (Ident "w")) (Var (Ident "x"))

Ident "w" := Val 3
Ident "x" := Add (Val 3) (Val 3)
Ident "y" := Add (Val 3) (Add (Val 1) (Var (Ident "x")))
Ident "x" := Add (Var (Ident "y")) (Var (Ident "y"))
Ident "y" := Add (Val 3) (Var (Ident "x"))

Ident "w" := Val 3
Ident "x" := Val 6
Ident "y" := Add (Val 3) (Add (Val 1) (Val 6))
Ident "x" := Add (Var (Ident "y")) (Var (Ident "y"))
Ident "y" := Add (Val 3) (Var (Ident "x"))

Ident "w" := Val 3
Ident "x" := Val 6
Ident "y" := Val 10
Ident "x" := Add (Val 10) (Val 10)
Ident "y" := Add (Val 3) (Var (Ident "x"))

Ident "w" := Val 3
Ident "x" := Val 6
Ident "y" := Val 10
Ident "x" := Val 20
Ident "y" := Add (Val 3) (Val 20)

Ident "w" := Val 3
Ident "x" := Val 6
Ident "y" := Val 10
Ident "x" := Val 20
Ident "y" := Val 23

Fin.

newtype Ident

Constructors

Ident String

Instances

Eq Ident

Methods

(==) :: Ident -> Ident -> Bool

(/=) :: Ident -> Ident -> Bool

Data Ident

Methods

gfoldl :: (forall d b. Data d => c (d -> b) -> d -> c b) -> (forall g. g -> c g) -> Ident -> c Ident

gunfold :: (forall b r. Data b => c (b -> r) -> c r) -> (forall r. r -> c r) -> Constr -> c Ident

toConstr :: Ident -> Constr

dataTypeOf :: Ident -> DataType

dataCast1 :: Typeable (* -> *) t => (forall d. Data d => c (t d)) -> Maybe (c Ident)

dataCast2 :: Typeable (* -> * -> *) t => (forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c Ident)

gmapT :: (forall b. Data b => b -> b) -> Ident -> Ident

gmapQl :: (r -> r' -> r) -> r -> (forall d. Data d => d -> r') -> Ident -> r

gmapQr :: (r' -> r -> r) -> r -> (forall d. Data d => d -> r') -> Ident -> r

gmapQ :: (forall d. Data d => d -> u) -> Ident -> [u]

gmapQi :: Int -> (forall d. Data d => d -> u) -> Ident -> u

gmapM :: Monad m => (forall d. Data d => d -> m d) -> Ident -> m Ident

gmapMp :: MonadPlus m => (forall d. Data d => d -> m d) -> Ident -> m Ident

gmapMo :: MonadPlus m => (forall d. Data d => d -> m d) -> Ident -> m Ident

Ord Ident

Methods

compare :: Ident -> Ident -> Ordering

(<) :: Ident -> Ident -> Bool

(<=) :: Ident -> Ident -> Bool

(>) :: Ident -> Ident -> Bool

(>=) :: Ident -> Ident -> Bool

max :: Ident -> Ident -> Ident

min :: Ident -> Ident -> Ident

Show Ident

Methods

showsPrec :: Int -> Ident -> ShowS

show :: Ident -> String

showList :: [Ident] -> ShowS

data Expr

Constructors

Val Int

Add Expr Expr

Var Ident

Instances

Eq Expr

Methods

(==) :: Expr -> Expr -> Bool

(/=) :: Expr -> Expr -> Bool

Data Expr

Methods

gfoldl :: (forall d b. Data d => c (d -> b) -> d -> c b) -> (forall g. g -> c g) -> Expr -> c Expr

gunfold :: (forall b r. Data b => c (b -> r) -> c r) -> (forall r. r -> c r) -> Constr -> c Expr

toConstr :: Expr -> Constr

dataTypeOf :: Expr -> DataType

dataCast1 :: Typeable (* -> *) t => (forall d. Data d => c (t d)) -> Maybe (c Expr)

dataCast2 :: Typeable (* -> * -> *) t => (forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c Expr)

gmapT :: (forall b. Data b => b -> b) -> Expr -> Expr

gmapQl :: (r -> r' -> r) -> r -> (forall d. Data d => d -> r') -> Expr -> r

gmapQr :: (r' -> r -> r) -> r -> (forall d. Data d => d -> r') -> Expr -> r

gmapQ :: (forall d. Data d => d -> u) -> Expr -> [u]

gmapQi :: Int -> (forall d. Data d => d -> u) -> Expr -> u

gmapM :: Monad m => (forall d. Data d => d -> m d) -> Expr -> m Expr

gmapMp :: MonadPlus m => (forall d. Data d => d -> m d) -> Expr -> m Expr

gmapMo :: MonadPlus m => (forall d. Data d => d -> m d) -> Expr -> m Expr

Show Expr

Methods

showsPrec :: Int -> Expr -> ShowS

show :: Expr -> String

showList :: [Expr] -> ShowS

constFold :: Expr -> Expr

substitute :: Map Ident Int -> Expr -> Expr

data Stmt

Instances

Show Stmt

Methods

showsPrec :: Int -> Stmt -> ShowS

show :: Stmt -> String

showList :: [Stmt] -> ShowS

constProp :: Map Ident Int -> Stmt -> (Map Ident Int, Stmt)

constFold' :: Expr -> Edit Expr

constFoldPass :: [Stmt] -> Edit [Stmt]

substitute' :: Map Ident Int -> Expr -> Edit Expr

constProp' :: Map Ident Int -> Stmt -> (Map Ident Int, Edit Stmt)

constPropPass :: [Stmt] -> Edit [Stmt]

constFoldAndPropPass :: [Stmt] -> [Stmt]