@@ -15,7 +15,7 @@ module BodyBuilder (
1515
1616import AST
1717import Debug.Trace
18- import Snippets ( boolType , intType , primMove )
18+ import Snippets ( boolType , intType , primMove , primCast )
1919import Util
2020import Config (minimumSwitchCases , wordSize )
2121import Options (LogSelection (BodyBuilder ))
@@ -30,14 +30,16 @@ import Data.Bits
3030import Data.Function
3131import Data.Functor
3232import Control.Monad
33- import Control.Monad.Extra (whenJust , whenM , fold1M )
33+ import Control.Monad.Extra (whenJust , whenM , fold1M , ifM )
3434import Control.Monad.Trans (lift )
3535import Control.Monad.Trans.State
3636import AST (simpleShowMap )
3737import Data.Foldable.Extra (foldlM , foldrM )
3838import Control.Monad.Trans.Maybe (MaybeT (runMaybeT , MaybeT ))
3939import Control.Applicative
4040import qualified Control.Applicative as App
41+ import qualified Data.Foldable as List
42+ import qualified Data.Foldable.Extra as List
4143
4244
4345----------------------------------------------------------------
@@ -468,7 +470,7 @@ instr' prim@(PrimForeign "lpvm" "cast" []
468470 [from, to@ ArgVar {argVarName= var, argVarFlow= flow}]) pos = do
469471 logBuild $ " Expanding cast(" ++ show from ++ " , " ++ show to ++ " )"
470472 unless (argFlowDirection from == FlowIn && flow == FlowOut ) $
471- shouldnt " cast instruction with wrong flow"
473+ shouldnt $ " cast instruction with wrong flow" ++ show prim
472474 if argType from == argType to
473475 then instr' (PrimForeign " llvm" " move" [] [from, to]) pos
474476 else ordinaryInstr prim pos
@@ -1605,16 +1607,31 @@ markIfLastUse _ arg = arg
16051607-- Yields a ProcBody with the fork, or NoFork if all branches can be merged with no substitutions
16061608rebuildFork :: PrimVarName -> TypeSpec -> Bool -> [(Integer , ProcBody )] -> Maybe ProcBody -> BkwdBuilder ProcBody
16071609rebuildFork var ty lastUse brs dflt = do
1608- mbMerged <- runMaybeT $ mergeBranches brs dflt $ Factors Map. empty Map. empty
1610+ mbMerged <- runMaybeT $ mergeBranches brs dflt $ Factors Map. empty Map. empty
16091611 case mbMerged of
16101612 Just (Factors {factored= factored}, (mergedBody, dflt'))
1611- | Map. null factored ->
1612- bkwdFreshVar <&> \ tmp -> guardedMergedFork tmp var ty (maximum (fst <$> brs)) mergedBody dflt'
1613- | otherwise ->
1614- return $ ProcBody [] $ MergedFork var ty lastUse (flattenFactored <$> Map. assocs factored) mergedBody dflt'
1613+ -> do
1614+ (table, finalBody, lastUse') <- constructTable factored mergedBody
1615+ if List. null table
1616+ then bkwdFreshVar <&> \ tmp -> guardedMergedFork tmp var ty (maximum (fst <$> brs)) finalBody dflt'
1617+ else return $ ProcBody [] $ MergedFork var ty lastUse' table finalBody dflt'
16151618 _ -> return $ ProcBody [] $ PrimFork var ty lastUse brs dflt
16161619 where
1617- flattenFactored (var,(ty,vals)) = (var, ty, vals ++ replicate (length brs - length vals) (last vals))
1620+ -- Ensure that all factored entries have the same length by padding with the last value if needed.
1621+ -- The last value is used because the branches are merged in reverse.
1622+ -- Also remove all factored variables that are of the form [0..].
1623+ constructTable table body = do
1624+ (table', casts, renamed) <- List. foldrM
1625+ (\ (tmp, (tmpTy, vals)) (table, casts, renamed) ->
1626+ if Maybe. mapMaybe argIntegerValue vals == [0 .. genericLength vals- 1 ]
1627+ then
1628+ -- avoid looking up the type rep if theyre already equal
1629+ ifM (return (ty == tmpTy) <||> (liftA2 (==) `on` lift . lookupTypeRepresentation) ty tmpTy)
1630+ (return (table, casts, (tmp,var): renamed))
1631+ (return (table, Unplaced (primCast (ArgVar tmp tmpTy FlowOut Ordinary False ) (ArgVar var ty FlowIn Ordinary (lastUse && List. null casts))): casts, renamed))
1632+ else return ((tmp, tmpTy, vals ++ replicate (length brs - length vals) (last vals)): table, casts, renamed)
1633+ ) ([] , [] , [] ) $ Map. assocs table
1634+ return (table', prependToBody casts $ renameProcBody (Map. fromList renamed) body, List. null casts)
16181635
16191636-- | When merging, we replace constants in both things at the same place with variables, tracking what has been merged in a map.
16201637-- When we merge again with something else, we can consult the map to see if a now variable was previously a constant.
@@ -1623,7 +1640,7 @@ rebuildFork var ty lastUse brs dflt = do
16231640-- LHS to a name from the RHS
16241641data Factors = Factors {
16251642 factored :: Map PrimVarName (TypeSpec , [PrimArg ]), -- ^ The variables and corresponding constants that have been factored
1626- renamed :: Map PrimVarName PrimVarName -- ^ The variables to rename on the RHS
1643+ renamed :: VarSubstitution -- ^ The variables to rename on the RHS
16271644} deriving (Show )
16281645
16291646type FactoredMerge a = Factors -> MaybeT BkwdBuilder (Factors , a )
@@ -1636,7 +1653,7 @@ combineMerged joiner mergeA mergeB vars = do
16361653
16371654-- | Try to merge all branches, with Nothing if the branches cannot be merged
16381655mergeBranches :: [(Integer , ProcBody )] -> Maybe ProcBody -> FactoredMerge (ProcBody , Maybe ProcBody )
1639- mergeBranches brs@ (_ : _) dflt vars = (do
1656+ mergeBranches brs dflt vars = (do
16401657 lift $ logBkwd $ " Trying to merge " ++ intercalate " \n " (show <$> brs) ++ " \n with default: " ++ show dflt
16411658 (vars', body') <- foldrM mergeBodies' (vars, last bodies') $ init bodies'
16421659 mbMergedDflt <- join <$> optional (forM dflt $ flip mergeBodies' (vars', body'))
@@ -1646,7 +1663,7 @@ mergeBranches brs@(_:_) dflt vars = (do
16461663 -> return (vars'', (dflt', Nothing ))
16471664 _ ->
16481665 -- when dense, can merge all branches, keeping default
1649- if idxs == List. take ( length brs) [0 .. ]
1666+ if idxs == [0 .. genericLength brs - 1 ]
16501667 then return (vars', (body', dflt))
16511668 else App. empty
16521669 lift $ logBkwd $ " Merged into " ++ show mergedAll
@@ -1655,7 +1672,6 @@ mergeBranches brs@(_:_) dflt vars = (do
16551672 where
16561673 (idxs, bodies') = unzip brs
16571674 mergeBodies' b (v, a) = mergeBodies a b v{renamed= renamed vars}
1658- mergeBranches _ _ _ = shouldnt " "
16591675
16601676-- | Try to merge two ProcBodies, merging the prims pairwise
16611677mergeBodies :: ProcBody -> ProcBody -> FactoredMerge ProcBody
@@ -1689,14 +1705,14 @@ mergeForks (PrimFork var0 ty0 final0 branches0 dflt0) (PrimFork var1 ty1 final1
16891705 | v0 < v1
16901706 = case dflt1 of
16911707 Nothing -> ((v0,b0): ) <$$> mergeIndexedBranches bs0 bbs1 vars
1692- Just dflt1' -> mergeIndexedBranches bbs0 ((v0,dflt1'): bbs1) vars
1708+ Just dflt1' -> mergeIndexedBranches bbs0 ((v0,dflt1'): bbs1) vars
16931709 | v0 > v1
16941710 = case dflt0 of
16951711 Nothing -> ((v0,b0): ) <$$> mergeIndexedBranches bs0 bbs1 vars
16961712 Just dflt0' -> mergeIndexedBranches ((v1,dflt0'): bbs0) bbs1 vars
16971713 | otherwise
1698- = combineMerged ((:) . (v0,))
1699- (mergeBodies b0 b1)
1714+ = combineMerged ((:) . (v0,))
1715+ (mergeBodies b0 b1)
17001716 (mergeIndexedBranches bs0 bs1)
17011717 vars
17021718 mergeIndexedBranches _ _ _ = App. empty
@@ -1705,9 +1721,9 @@ mergeForks (MergedFork var0 ty0 final0 table0 branch0 dflt0) (MergedFork var1 ty
17051721 if tyRep0 == tyRep1 && var0 == rename renamed var1 && final0 == final1
17061722 then combineMerged (uncurry . MergedFork var0 ty0 final0)
17071723 (mergeTables table0 table1)
1708- (combineMerged (,)
1709- (mergeBodies branch0 branch1)
1710- (mergeMaybeBodies dflt0 dflt1))
1724+ (combineMerged (,)
1725+ (mergeBodies branch0 branch1)
1726+ (mergeMaybeBodies dflt0 dflt1))
17111727 vars
17121728 else App. empty
17131729 where
@@ -1787,7 +1803,8 @@ mergeLists f as0 as1 vars =
17871803 foldrM (\ (a0, a1) (vars', as') -> MaybeT (mapSnd (: as') <$$> runMaybeT (f a0 a1 vars'))) (vars, [] )
17881804 $ zip as0 as1
17891805
1790- renameProcBody :: Map PrimVarName PrimVarName -> ProcBody -> ProcBody
1806+ -- \ Rename all variables in a proc body, under a given VarSubsitution
1807+ renameProcBody :: VarSubstitution -> ProcBody -> ProcBody
17911808renameProcBody vars (ProcBody prims fork) = ProcBody (contentApply (renamePrim vars) <$> prims) (renameFork vars fork)
17921809 where
17931810 renamePrim vars prim =
@@ -1804,7 +1821,8 @@ renameProcBody vars (ProcBody prims fork) = ProcBody (contentApply (renamePrim v
18041821 renamePrimArg vars (ArgClosure pspec args ty) = ArgClosure pspec (renamePrimArg vars <$> args) ty
18051822 renamePrimArg _ arg = arg
18061823
1807- rename :: Map PrimVarName PrimVarName -> PrimVarName -> PrimVarName
1824+ -- | Rename a variable, defaulting to itself if there is no renaming
1825+ rename :: VarSubstitution -> PrimVarName -> PrimVarName
18081826rename vars var = fromMaybe var $ Map. lookup var vars
18091827
18101828
0 commit comments