content: module DDC.Core.Llvm.Metadata.Tbaa
         ( MDSuper(..)
         , deriveMD
         , annot
         , lookup, lookups )
  where
  import DDC.Llvm.Syntax.Metadata
  import DDC.Core.Collect
  import DDC.Type.Env                     (KindEnv)
  import DDC.Type.Exp.Simple
  import DDC.Core.Llvm.Metadata.Graph
  import DDC.Core.Llvm.Convert.Base
  import qualified DDC.Type.Env           as Env
  import qualified DDC.Core.Salt          as A
  import qualified DDC.Llvm.Syntax        as V
  
  import Prelude                          hiding (lookup)
  import Control.Applicative
  import Control.Monad
  import Data.Maybe
  import Data.Map                         (Map)
  import Data.List                        hiding (lookup)
  import qualified Data.Map               as Map
  import qualified Data.Set               as Set
  import qualified Data.Text              as T
  
  
  -- Metadata management --------------------------------------------------------
  -- | Metadata for a supercombinator.
  data MDSuper
          = MDSuper
            { -- Map bound regions to metadata nodes for attaching metadata
              --    to relevant instructions.
              nameMap     :: MDEnv
  
              -- Metadata nodes, to be pretty-printed with the module as
              --    declarations. e.g. "1 = !{ metadata "id", !parent, !i11}
            , decls       :: [MDecl]
            } deriving Show
  
  
  -- | Map region variables to relevant metadata
  --      need the whole declaration for tags, e.g "!tbaa", "!debug"
  type MDEnv = Map (Bound A.Name) [MDecl]
  
  emptyDict :: MDEnv
  emptyDict = Map.empty
  
  extendDict :: (Bound A.Name, MDecl) -> MDEnv -> MDEnv
  extendDict (u, n) e | Map.member u e = Map.adjust (n:) u e
                      | otherwise      = Map.insert u [n] e
  
  
  -- | Lookup the metadata for a name, from the metadata tree attached
  --   to a supecombinator.
  lookup :: Bound A.Name -> MDSuper -> Maybe [MDecl]
  lookup u mdsup = Map.lookup u (nameMap mdsup)
  
  
  -- | Like `lookup` but lookup metadata for several names at once.
  lookups :: [Bound A.Name] -> MDSuper -> [Maybe [MDecl]]
  lookups us mdsup = map (flip lookup mdsup) us
  
  
  -- | Generate tbaa metadata for a top-level Salt supercombinator.
  deriveMD
        :: (BindStruct A.Exp A.Name)
        => String                 -- ^ Sanitized name of super
        -> A.Exp                  -- ^ Super to derive from
        -> ConvertM MDSuper       -- ^ Metadata encoding witness information
  
  deriveMD nTop xx
    = let
          regs                 = collectRegsB xx
          (constwits, diswits) = partitionWits $ collectWitsB xx
          arel                 = constructARel   diswits
          domain               = constructANodes regs constwits
          mdDG                 = orientUG $ UG (domain, arel)
          mdTrees              = partitionDG mdDG
      in  foldM (buildMDTree nTop) (MDSuper emptyDict []) mdTrees
  
  
  buildMDTree :: String -> MDSuper -> Tree ANode -> ConvertM MDSuper
  buildMDTree nTop sup tree
   = let tree' = anchor ARoot tree
     in  bfBuild nTop tree' Nothing sup ARoot
  
  
  bfBuild :: String -> Tree ANode -> Maybe MRef -> MDSuper -> ANode -> ConvertM MDSuper
  bfBuild nTop tree parent sup node
   = case parent of
          Nothing        -> do name <- freshRootName nTop
                               bf Nothing $ tbaaRoot name
  
          Just parentRef -> do name <- freshNodeName nTop (regionU node)
                               bf (Just $ regionU node) $ tbaaNode name parentRef (isConst node)
     where bf u md
            = do ref          <- liftM MRef $ newUnique
                 let sup'     =  declare u ref md sup
                 let children =  sources node tree
                 foldM (bfBuild nTop tree (Just ref)) sup' children
           declare u r m s
            = let decl = MDecl r m
              in  case u of Nothing -> s { decls   = decl:(decls s) }
                            Just u' -> s { nameMap = extendDict (u',decl) $ nameMap s
                                         , decls   = decl:(decls s) }
  
  
  freshNodeName :: String -> Bound A.Name -> ConvertM String
  freshNodeName q (UName nm)
      | Just n <- A.takeNameVar nm
      = return $ q ++ "_" ++ (T.unpack n)
  
  freshNodeName q _
      = liftA (\i -> q ++ "_" ++ (show i)) newUnique
  
  freshRootName :: String -> ConvertM String
  freshRootName qualify = liftA (\i -> qualify ++ "_ROOT_" ++ (show i)) newUnique
  
  
  -- | Attach relevant metadata to instructions
  annot :: (BindStruct (c A.Name) A.Name)
        => KindEnv A.Name
        -> MDSuper        -- ^ Metadata
        -> [c A.Name]     -- ^ Things to lookup for Meta data.
        -> V.Instr        -- ^ Instruction to annotate
        -> V.AnnotInstr
  
  annot kenv mdsup xs ins
   = let regions  = concatMap (collectRegsU kenv) xs
         mdecls   = concat $ catMaybes $ lookups regions mdsup
         annotate' ms is
           = case is of
                  V.ILoad{}  -> V.AnnotInstr is ms
                  V.IStore{} -> V.AnnotInstr is ms
                  _          -> V.AnnotInstr is []
     in  annotate' mdecls ins
  
  
  -- Alias relation -------------------------------------------------------------
  -- | A node in the alias graphs, representing a region
  data ANode  = ANode { regionU :: RegBound
                      , isConst :: Bool }
              | ARoot
                deriving (Show, Eq, Ord)
  
  
  -- | Make nodes from regions
  constructANodes :: [RegBound] -> [WitType] -> [ANode]
  constructANodes regs constwits
   = let isConstR r = or $ map (flip isConstWFor r) constwits
         mkANode r  = ANode r (isConstR r)
     in  map mkANode regs
  
  
  -- | Encode the `alias` relation defined on the set regions
  --      * reflexitivity is taken as implicit
  --        (important for generating DAG, and safe since LLVM already assumes
  --         reflexitivity)
  --      * symmetry is made explicit
  --      note that `alias` is non-transitive.
  --
  constructARel :: [WitType] -> Rel ANode
  constructARel diswits = alias
    where alias n1 n2
            | n1 == n2  = False
            | otherwise = not $ or
                        $ map (flip isDistinctWFor (regionU n1, regionU n2)) diswits
  
  
  -- Collecting bounds ----------------------------------------------------------
  type RegBound  = Bound A.Name
  type WitType   = Type A.Name
  
  
  isConstW :: WitType -> Bool
  isConstW t  = isConstWitType t
  
  
  isConstWFor :: WitType -> RegBound -> Bool
  isConstWFor t r
    | _ : args <- takeTApps t
    = and [isConstWitType t, elem (TVar r) args]
    | otherwise = False
  
  
  isDistinctW :: WitType-> Bool
  isDistinctW tw
    | tc : _ <- takeTApps tw = isDistinctWitType tc
    | otherwise              = False
  
  
  isDistinctWFor :: WitType -> (RegBound, RegBound) -> Bool
  isDistinctWFor t (r1,r2)
    | tc : args <- takeTApps t
    = and [isDistinctWitType tc, (TVar r1) `elem` args, (TVar r2) `elem` args]
    | otherwise = False
  
  
  -- | Divide a set of witnesses to a set of Const wits and a set of Distinct wits
  partitionWits :: [WitType] -> ([WitType], [WitType])
  partitionWits ws
    = partition isConstW
    $ filter    (liftA2 (||) isConstW isDistinctW) ws
  
  
  -- | Collect region bounds
  collectRegsU :: (BindStruct (c A.Name) A.Name) => KindEnv A.Name -> c A.Name -> [RegBound]
  collectRegsU kenv cc
   = let isReg u = case Env.lookup u kenv of
                        Just t | isRegionKind t -> True
                        _                       -> False
     in  filter isReg $ Set.toList (collectBound cc)
  
  
  -- | Collect region bindings
  collectRegsB :: (BindStruct (c A.Name) A.Name) => c A.Name -> [RegBound]
  collectRegsB cc
   = let isBindReg b
           = case b of
                  BName n t | isRegionKind t -> Just (UName n)
                  _                          -> Nothing
         bindRegs = map (isBindReg) $ fst (collectBinds cc)
     in  catMaybes bindRegs
  
  
  -- | Collect witness bindings together with their types (for convinience)
  collectWitsB :: (BindStruct (c A.Name) A.Name) => c A.Name -> [WitType]
  collectWitsB cc
   = let isBindWit b
          = let t = typeOfBind b
            in  if isWitnessType t then Just t else Nothing
         bindWits  = map (isBindWit) $ snd (collectBinds cc)
     in  catMaybes bindWits
  
  