Using Tardis to turn generated names into nicely readable ones

ReadableGeneratedNames.hs

{-# LANGUAGE DataKinds, KindSignatures #-}
{-# LANGUAGE GADTs, StandaloneDeriving #-}
import Prelude hiding (mapM)
import Data.Traversable (mapM)

import Control.Monad.Tardis
import Control.Monad.Free
import Data.Stream (Stream(..))
import qualified Data.Stream as Stream
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Monoid
import Control.Arrow (first, second)

-- Our imaginary compiler has a single pass...
data Phase = Parsed | Processed
           deriving Show

-- ... that introduces new names into our program.
data Name (a :: Phase) where
    Given :: String -> Name a
    Generated :: Int -> Name Processed
deriving instance Show (Name a)

-- We model the 'program' as just a tree of names.
type Tree a = Free [] (Name a)

-- An example 'program' that has already been processed. It contains
-- the name "x1" that was introduces by the programmer, and two
-- references to a single, generated name.
test :: Tree Processed
test = Free [v2, Free [v2, v1]]
  where
    v1 = Pure (Given "x1")
    v2 = Pure (Generated 1)

-- We want to make programs like this nicely readable by assigning
-- friendly names (like "x1") to the variables introduced by the
-- compiler. Of course, we don't want clashes between user-defined and
-- compiler-defined names.

-- We collect the user-defined names and send them backwards in
-- time. The forwards-flowing state is a stream of fresh names, and a
-- mapping of generated names already encountered.
type Readable a = Tardis (Set String) (Stream String, Map Int String) a

-- Getting a fresh name is easy.
fresh :: Readable String
fresh = do
    ~(Cons s ss) <- getsPast fst
    modifyForwards . first $ const ss
    return s

-- Making a name readable is a simple matter of three possible cases:
-- sending it back in time, looking up its already assigned name, or
-- assigning a fresh name to it.
readableName :: Name a -> Readable (Name Parsed)
readableName (Given s) = do
    modifyBackwards $ Set.insert s
    return $ Given s
readableName (Generated x) = fmap Given $ do
    lookup <- getsPast $ Map.lookup x . snd
    case lookup of
        Just name -> return name
        Nothing -> do
            name <- fresh
            modifyForwards . second $ Map.insert x name
            return name

-- Programs are made readable by traversing over them
readableTree :: Tree a -> Readable (Tree Parsed)
readableTree = mapM readableName

-- And here we tie the knot via `runTardis`
readable :: Tree a -> Tree Parsed
readable t = t'
  where
    (t', (usedNames, _)) = runTardis (readableTree t) (mempty, (names, mempty))
    names0 = fmap (\i -> 'x':show i) $ Stream.iterate succ 1
    names = Stream.filter (not . (`Set.member` usedNames)) names0