Inline top level aliases at their call sites

src/Pukeko.hs

@@ -10,6 +10,7 @@ import           Pukeko.Error
 
 import           Pukeko.Language.AST.Std        (Module (..))
 import           Pukeko.Language.AST.Stage      (Typed)
+import qualified Pukeko.Language.AliasInliner   as AliasInliner
 import qualified Pukeko.Language.CoreCompiler   as CoreCompiler
 import qualified Pukeko.Language.DeadCode       as DeadCode
 import qualified Pukeko.Language.EtaReducer     as EtaReducer
@@ -18,6 +19,7 @@ import qualified Pukeko.Language.KindChecker    as KindChecker
 import qualified Pukeko.Language.LambdaLifter   as LambdaLifter
 import qualified Pukeko.Language.Parser         as Parser
 import qualified Pukeko.Language.PatternMatcher as PatternMatcher
+import qualified Pukeko.Language.Prettifier     as Prettifier
 import qualified Pukeko.Language.Inferencer     as Inferencer
 import qualified Pukeko.Language.TypeChecker    as TypeChecker
 import qualified Pukeko.Language.TypeResolver   as TypeResolver
@@ -42,6 +44,8 @@ compileToCore unsafe module_pu = do
                >>= typeChecked PatternMatcher.compileModule
                >>= typeChecked (pure . LambdaLifter.liftModule)
                >>= typeChecked (pure . EtaReducer.reduceModule)
+               >>= typeChecked (pure . AliasInliner.inlineModule)
                >>= typeChecked (pure . DeadCode.cleanModule)
+               >>= typeChecked (pure . Prettifier.prettifyModule)
   let module_cc = CoreCompiler.compileModule module_ll
   return (module_cc, module_ll, module_ti)

src/Pukeko/Language/AST/Std.hs

@@ -33,6 +33,9 @@ module Pukeko.Language.AST.Std
   , altnRhs
 
   , module2tops
+  , top2lhs
+  , top2rhs
+  , top2eval
   , defn2rhs
   , defn2dcon
   , expr2eval
@@ -201,11 +204,40 @@ module2tops ::
   Lens (Module st1) (Module st2) [TopLevel st1] [TopLevel st2]
 module2tops f (MkModule info tops) = MkModule info <$> f tops
 
+top2lhs ::
+  (HasTLTyp st ~ 'False, HasTLVal st ~ 'False) =>
+  IndexedLens' Pos (TopLevel st) Id.EVar
+top2lhs f = \case
+  TLDef d -> TLDef <$> defn2lhs (bind2evar f) d
+  TLSup w z vs t xs e -> (\z' -> TLSup w z' vs t xs e) <$> indexed f w z
+  TLAsm b s -> (\b' -> TLAsm b' s) <$> bind2evar f b
+
+top2rhs ::
+  (Applicative f) =>
+  (forall tv ev. Expr st tv ev -> f (Expr st tv ev)) -> TopLevel st -> f (TopLevel st)
+top2rhs f = \case
+  TLTyp w ds -> pure (TLTyp w ds)
+  TLVal w z t -> pure (TLVal w z t)
+  TLDef d -> TLDef <$> defn2rhs f d
+  TLSup w z vs t bs e -> TLSup w z vs t bs <$> f e
+  TLAsm b s -> pure (TLAsm b s)
+
+top2eval :: IndexedTraversal' Pos (TopLevel st) Id.EVar
+top2eval f = top2rhs (expr2eval f)
+
+defn2lhs :: Lens' (Defn st tv ev) (Bind st tv)
+defn2lhs = defnLhs
+
 defn2rhs ::
   (StageType st1 ~ StageType st2) =>
   Lens (Defn st1 tv ev1) (Defn st2 tv ev2) (Expr st1 tv ev1) (Expr st2 tv ev2)
 defn2rhs f (MkDefn b e) = MkDefn (retagBind b) <$> f e
 
+bind2evar ::
+  (StageType st1 ~ StageType st2) =>
+  IndexedLens Pos (Bind st1 tv) (Bind st2 tv) Id.EVar Id.EVar
+bind2evar f (MkBind w x t) = fmap (\x' -> MkBind w x' t) (indexed f w x)
+
 -- * Deep traversals
 type DConTraversal t =
   forall st1 st2 tv ev. (SameNodes st1 st2, SameTypes st1 st2) =>

src/Pukeko/Language/AliasInliner.hs

@@ -0,0 +1,71 @@
+module Pukeko.Language.AliasInliner
+  ( inlineModule
+  ) where
+
+import           Control.Lens      hiding (indices)
+import           Control.Monad.ST
+import qualified Data.Array        as A
+import qualified Data.Array.ST     as A
+import           Data.Foldable
+import qualified Data.Map          as Map
+import           Data.Maybe        (mapMaybe)
+import           Data.Monoid
+import qualified Data.Set          as Set
+
+import           Pukeko.Language.AST.Std
+import           Pukeko.Language.AST.Classes
+import qualified Pukeko.Language.Ident   as Id
+
+inlineModule :: Module st -> Module st
+inlineModule (MkModule info tops0) =
+  let ls = mapMaybe topLink tops0
+      uf = unionFind ls
+      tops1 = over (traverse . top2eval) (\x -> Map.findWithDefault x x uf) tops0
+  in  MkModule info tops1
+
+topLink :: TopLevel st -> Maybe (Id.EVar, Id.EVar)
+topLink = \case
+  TLDef (MkDefn b (EVal _ x)) -> Just (b^.lhs, x)
+  TLSup _ z vs _ xs (EVal _ x)
+    | null vs && null xs -> Just (z, x)
+  _ -> Nothing
+
+data UnionFind st a = UnionFind
+  { indices :: Map.Map a Int
+  , names   :: A.Array Int a
+  , links   :: A.STArray st Int Int
+  }
+
+newUnionFind :: Ord a => Set.Set a -> ST s (UnionFind s a)
+newUnionFind xs = do
+  let is = Map.fromList (zip (toList xs) [0..])
+  let bnds = (0, Map.size is - 1)
+  let ns = A.array bnds (map (\(x, i) -> (i, x)) (Map.toList is))
+  ls <- A.newListArray bnds [0..]
+  pure (UnionFind is ns ls)
+
+root :: UnionFind s a -> Int -> ST s Int
+root uf i = do
+  j <- A.readArray (links uf) i
+  if i /= j
+    then do
+      k <- A.readArray (links uf) j
+      A.writeArray (links uf) i k
+      root uf j
+    else
+      pure i
+
+join' :: UnionFind s a -> Int -> Int -> ST s ()
+join' uf i j = do
+  l <- root uf j
+  A.writeArray (links uf) i l
+
+join :: Ord a => UnionFind s a -> a -> a -> ST s ()
+join uf x y = join' uf (indices uf Map.! x) (indices uf Map.! y)
+
+unionFind :: Ord a => [(a, a)] -> Map.Map a a
+unionFind xys = runST $ do
+  let xs = foldMap (\(x, y) -> Set.singleton x <> Set.singleton y) xys
+  uf <- newUnionFind xs
+  for_ xys (uncurry (join uf))
+  traverse (\i -> (names uf A.!) <$> root uf i) (indices uf)

src/Pukeko/Language/DeadCode.hs

@@ -10,7 +10,6 @@ import qualified Data.Map      as Map
 import qualified Data.Set      as Set
 import qualified Data.Set.Lens as Set
 
-import           Pukeko.Language.AST.Classes
 import           Pukeko.Language.AST.Std
 import           Pukeko.Language.AST.ModuleInfo (info2funs)
 import qualified Pukeko.Language.AST.Stage      as St
@@ -20,26 +19,12 @@ type ElimStage st = (St.HasTLTyp st ~ 'False, St.HasTLVal st ~ 'False)
 
 cleanModule :: (ElimStage st) => Module st -> Module st
 cleanModule (MkModule info0 tops0) =
-  let (g, out, in_) = G.graphFromEdges $ map (\t -> (t, topLevelLhs t, deps t)) tops0
+  let (g, out, in_) = G.graphFromEdges $ map (\t -> (t, t^.top2lhs, deps t)) tops0
       reach = Set.fromList
               $ map (view _2 . out) $ maybe [] (G.reachable g) (in_ Id.main)
       keep = (`Set.member` reach)
       info1 = over info2funs (Map.filterWithKey (const . keep)) info0
-      tops1 = filter (keep . topLevelLhs) tops0
+      tops1 = filter (keep . view top2lhs) tops0
   in  MkModule info1 tops1
   where
-    deps = Set.toList . Set.setOf (\f -> topLevel2expr (expr2eval f))
-
-topLevelLhs :: (ElimStage st) => TopLevel st -> Id.EVar
-topLevelLhs = \case
-  TLDef (MkDefn b _) -> b^.lhs
-  TLSup _ z _ _ _ _  -> z
-  TLAsm b _          -> b^.lhs
-
-topLevel2expr ::
-  (ElimStage st, Applicative f) =>
-  (forall tv ev. Expr st tv ev -> f (Expr st tv ev)) -> TopLevel st -> f (TopLevel st)
-topLevel2expr f = \case
-  TLDef d -> TLDef <$> defn2rhs f d
-  TLSup w z vs t bs e -> TLSup w z vs t bs <$> f e
-  TLAsm b s -> pure (TLAsm b s)
+    deps = Set.toList . Set.setOf top2eval

src/Pukeko/Language/Ident.hs

@@ -5,6 +5,7 @@ module Pukeko.Language.Ident
   , main
   , isVar
   , isOp
+  , stripPart
   , freshEVars
   , mangled
   , TVar
@@ -45,6 +46,9 @@ isVar _      = False
 isOp  Op{}   = True
 isOp  _      = False
 
+stripPart :: EVar -> EVar
+stripPart x = x{_part = Nothing}
+
 freshEVars :: String -> EVar -> [EVar]
 freshEVars comp var = map (\n -> var{_part = Just (comp, n)}) [1 ..]
 

src/Pukeko/Language/Prettifier.hs

@@ -0,0 +1,34 @@
+module Pukeko.Language.Prettifier
+  ( prettifyModule
+  ) where
+
+import           Control.Lens
+import qualified Data.Map      as Map
+import qualified Data.Set      as Set
+import qualified Data.Set.Lens as Set
+
+import           Pukeko.Language.AST.Std
+import           Pukeko.Language.AST.Stage
+import           Pukeko.Language.AST.ModuleInfo (info2funs)
+import qualified Pukeko.Language.Ident          as Id
+
+prettifyModule :: (HasTLTyp st ~ 'False, HasTLVal st ~ 'False) => Module st -> Module st
+prettifyModule (MkModule info0 tops0) =
+  let xs = Set.setOf (traverse . top2lhs) tops0
+      mp = cluster xs
+      rename x = Map.findWithDefault x x mp
+      info1 = over info2funs (Map.mapKeys rename) info0
+      tops1 = over (traverse . top2lhs) rename tops0
+      tops2 = over (traverse . top2eval) rename tops1
+  in  MkModule info1 tops2
+
+cluster :: Set.Set Id.EVar -> Map.Map Id.EVar Id.EVar
+cluster xs0 =
+  let mp =
+        foldl
+          (\acc x -> Map.insertWith Set.union (Id.stripPart x) (Set.singleton x) acc)
+          Map.empty xs0
+      f y xs1
+        | not (y `Set.member` xs1) = Map.singleton (Set.findMax xs1) y
+        | otherwise = Map.empty
+  in  ifoldMap f mp

test/Makefile

@@ -24,6 +24,7 @@ SOURCES     = monad_io.pu \
 							rmq.pu \
 							fix.pu \
 							lambdalift.pu \
+              aliasinline.pu \
 							sort_gen.pu \
 							rmq_gen.pu
 LAMBDAS     = $(SOURCES:.pu=.ll)

test/aliasinline.asm

@@ -0,0 +1,24 @@
+g_declare_cafs gm$cons_0_0, main
+g_declare_main main
+
+g_globstart gm$cons_0_0, 0
+g_updcons 0, 0, 1
+g_return
+
+g_globstart gm$return, 2
+g_updcons 0, 2, 1
+g_return
+
+g_globstart h, 1
+g_update 1
+g_unwind
+
+g_globstart main, 0
+g_pushglobal gm$cons_0_0, 0
+g_pushglobal h, 1
+g_mkap 1
+g_pushglobal h, 1
+g_mkap 1
+g_pushglobal gm$return, 2
+g_updap 1, 1
+g_unwind

test/aliasinline.ll

@@ -0,0 +1,3 @@
+external return : ∀a. a -> IO a = "return"
+let h : ∀a. a -> a = fun @a -> fun (u : a) -> u
+let main : IO Unit = return @Unit (h @Unit (h @Unit Unit))

test/aliasinline.pu

@@ -0,0 +1,14 @@
+val h : c -> c
+let h z = (fun u -> u) z
+
+val g : b -> b
+let g y = (fun z -> h z) y
+
+val f1 : a -> a
+let f1 = h
+
+val f2 : a -> a
+let f2 x = g x
+
+val main : IO Unit
+let main = return (f1 (f2 Unit))