<commit_before>module Data.LLVM.Analysis.PointsTo.Andersen (
  -- * Types
  Andersen,
  -- * Constructor
  runPointsToAnalysis,
  -- * Debugging aids
  viewPointsToGraph,
  ) where

import Data.Graph.Inductive hiding ( Gr, (><) )
import Data.GraphViz
import Data.IntMap ( IntMap )
import qualified Data.IntMap as IM
import Data.List ( foldl' )
import Data.Maybe ( mapMaybe )
import Data.Set ( Set )
import qualified Data.Set as S
import Data.Sequence ( Seq, (><), ViewL(..), viewl )
import qualified Data.Sequence as Seq

import Data.LLVM.Analysis.PointsTo
import Data.LLVM.Internal.PatriciaTree
import Data.LLVM.Types

import System.IO.Unsafe
import Text.Printf
import Debug.Trace
debug = flip trace

data NodeTag = Location Value
             -- | Argument Value
             -- | DirectCall Value
             -- | IndirectCall Value
             deriving (Ord, Eq, Show)

type PTGEdge = (Int, Int, ())
type PTGNode = (Int, NodeTag)
type PTG = Gr NodeTag ()
type Worklist = Seq Instruction
-- | Define a dependency graph.  The keys are the unique IDs of
-- locations.  The value for each ID is the set of instructions that
-- need to be re-processed when a new edge is added to the points-to graph
-- with that node ID as its source.
type DepGraph = IntMap (Set Instruction)

data Andersen = Andersen PTG

instance PointsToAnalysis Andersen where
  mayAlias (Andersen _) _ _ = True
  pointsTo = andersenPointsTo

andersenPointsTo :: (IsValue a) => Andersen -> a -> Set Value
andersenPointsTo (Andersen g) v =
  S.fromList $ map (unloc . lab g) (suc g (valueUniqueId v))
  where
    unloc (Just (Location l)) = l


-- | Run the points-to analysis and return an object that is an
-- instance of PointsToAnalysis, which can be used to query the
-- results.
runPointsToAnalysis :: Module -> Andersen
runPointsToAnalysis m = Andersen g
  where
    fs = moduleDefinedFunctions m
    blocks = concatMap functionBody fs
    insts = concatMap basicBlockInstructions blocks
    (globalLocations, globalEdges) = getGlobalLocations m
    argumentLocations = foldr extractArgs [] fs
    (localLocations, edgeInducers) = foldr extractLocations ([], []) insts
    allLocations = concat [globalLocations, argumentLocations, localLocations]
    -- The initial graph contains all locations in the program, along
    -- with edges induced by global initializers.
    graph0 = mkGraph allLocations globalEdges
    worklist0 = Seq.fromList edgeInducers
    g = saturate IM.empty worklist0 graph0

extractLocations :: Instruction
                    -> ([PTGNode], [Instruction])
                    -> ([PTGNode], [Instruction])
extractLocations i acc@(ptgNodes, insts) = case i of
  AllocaInst { instructionUniqueId = uid } ->
    ((uid, Location (Value i)) : ptgNodes, insts)
  StoreInst {} -> (ptgNodes, i : insts)
  CallInst {} -> (ptgNodes, i : insts)
  InvokeInst {} -> (ptgNodes, i : insts)
  _ -> acc

-- | Saturate the points-to graph using a worklist algorithm.
saturate :: DepGraph -> Worklist -> PTG -> PTG
saturate dg worklist g = case viewl worklist of
  EmptyL -> g
  itm :< rest -> case itm of
    StoreInst { storeValue = val, storeAddress = dest } ->
      case isPointerType (valueType val) of
        False -> saturate dg rest g
        True -> addStoreEdges dg itm val dest rest g
    CallInst {} -> saturate dg rest g
    InvokeInst {} -> saturate dg rest g
    _ -> error ("Unexpected instruction type: " ++ show itm)


-- | Handle adding edges induced by a @StoreInst@.
addStoreEdges :: DepGraph -> Instruction -> Value -> Value -> Worklist -> PTG -> PTG
addStoreEdges dg i val dest worklist g =
  -- Only update g and the worklist if there were new edges added.
  case null newEdges of
    False -> saturate dg2 worklist' g'
    -- IMPORTANT: Note that we need to propagate dependencies (the
    -- DepGraph) changes even if we didn't add any edges this time.
    -- Edges could be added later and we need to know all of the
    -- possible dependencies to come back.
    True -> saturate dg2 worklist g
  where
    accumUsedSrcs acc (_, src, _, _) = S.insert src acc
    (dg1, newTargets) = getLocationsReferencedByStore dg i val g
    (dg2, newSrcs) = getLocationsReferencedByStore dg1 i dest g
    newEdges = makeNewEdges g newSrcs newTargets
    usedSrcs = foldl accumUsedSrcs S.empty newEdges
    newWorklistItems = affectedInstructions usedSrcs dg2
    worklist' = worklist >< Seq.fromList newWorklistItems
    g' = foldl' (flip (&)) g newEdges

affectedInstructions :: Set Int -> DepGraph -> [Instruction]
affectedInstructions usedSrcs dg = S.toList instSet
  where
    instSet = S.fold findAffected S.empty usedSrcs
    findAffected nodeId acc = S.union acc (IM.findWithDefault S.empty nodeId dg)

-- | Determine which edges need to be added to the graph, based on the
-- set of discovered sources and targets.  Only new edges are
-- returned.
makeNewEdges :: PTG -> [Node] -> [Node] -> [Context NodeTag ()]
makeNewEdges g newSrcs newTargets =
  mapMaybe toContext newSrcs
  where
    notInGraph src tgt = notElem tgt (suc g src)
    -- | Create a new context for each src in the graph.  A context is
    -- only created if it adds new targets.  This is to make worklist
    -- management easier; if there are no new edges, the worklist is
    -- not updated.
    --
    -- Note that all edges for each source are added at once to the
    -- context.  Using separate contexts for each new edge would
    -- require changes later, otherwise some contexts could overwrite
    -- others, losing edges.
    toContext src =
      let targets = filter (notInGraph src) newTargets
          (adjIn, n, lbl, adjOut) = context g src
      in case targets of
        [] -> Nothing
        _ -> Just (adjIn, n, lbl, map (\t->((), t)) targets ++ adjOut)

-- | Given a @Value@ that is an operand of a @StoreInst@, find all of
-- the locations in the points-to graph that it refers to.  This means
-- that the function starts at the given value @v@ and _dereferences_
-- each @LoadInst@ in the @Value@.
--
-- As it progresses, the function also updates the DepGraph that
-- allows for reverse lookups (which instructions depend on the edges
-- coming from a given node).  This DepGraph is used to determine
-- which instructions need to be added to the Worklist due to
-- additional edges added to the points-to graph each iteration.  Note
-- that the DepGraph should be updated *before* it is used to compute
-- the list of affected instructions.
--
-- This function treats BitcastInsts as no-ops (i.e., it assumes the
-- types work out and only deals with memory references).
--
-- Algorithm: walk the chain of loads/geps/bitcasts to find the
-- location being loaded from.  Keep track of the number of
-- dereferences (loads).  Take that many steps (across all outgoing
-- edges) from the location in the points-to graph.
getLocationsReferencedByStore :: DepGraph
                                 -> Instruction
                                 -> Value
                                 -> PTG
                                 -> (DepGraph, [Node])
getLocationsReferencedByStore dg0 storeInst v g = getLocs v 0
  where
    getLocs :: Value -> Int -> (DepGraph, [Node])
    getLocs val derefCount = case valueContent val of
      -- This also needs to handle GEP instructions later, also
      -- (maybe) select, extractelement, and extractvalue.  This also
      -- needs to be careful around bitcasts of non-pointer types
      -- (integers turned into pointers, for example)
      InstructionC (LoadInst { loadAddress = addr }) ->
        getLocs addr (derefCount + 1)
      InstructionC (BitcastInst { castedValue = cv }) ->
        getLocs cv derefCount
      -- In this fallback case, @val@ should be a node in the
      -- Points-to graph.  Collect everything @derefCount@ steps from
      -- it and return that, along with the updated DepGraph.  The
      -- depgraph needs to be updated with an entry for each non-leaf
      -- node.
      _ -> collectLocationNodes g dg0 storeInst val derefCount

collectLocationNodes :: PTG -> DepGraph -> Instruction -> Value -> Int -> (DepGraph, [Node])
collectLocationNodes g dg0 storeInst val derefCount =
  collect dg0 derefCount [valueUniqueId val]
  where
    addDep dg n = IM.insertWith S.union n (S.singleton storeInst) dg
    collect dg 0 resultNodes = (dg, resultNodes)
    collect dg remainingHops startNodes =
      let nextNodes = concatMap (suc g) startNodes
          dg' = foldl' addDep dg startNodes
      in collect dg' (remainingHops - 1) nextNodes


-- | This only re-allocates the small part of the list each iteration,
-- so should remain efficient.<commit_msg>Stub out function call support<commit_after>module Data.LLVM.Analysis.PointsTo.Andersen (
  -- * Types
  Andersen,
  -- * Constructor
  runPointsToAnalysis,
  -- * Debugging aids
  viewPointsToGraph,
  ) where

import Data.Graph.Inductive hiding ( Gr, (><) )
import Data.GraphViz
import Data.IntMap ( IntMap )
import qualified Data.IntMap as IM
import Data.List ( foldl', mapAccumR, transpose )
import Data.Maybe ( mapMaybe )
import Data.Set ( Set )
import qualified Data.Set as S
import Data.Sequence ( Seq, (><), ViewL(..), viewl )
import qualified Data.Sequence as Seq

import Data.LLVM.Analysis.PointsTo
import Data.LLVM.Internal.PatriciaTree
import Data.LLVM.Types

import System.IO.Unsafe
import Text.Printf
import Debug.Trace
debug = flip trace

data NodeTag = Location Value
             -- | Argument Value
             -- | DirectCall Value
             -- | IndirectCall Value
             deriving (Ord, Eq, Show)

type PTGEdge = (Int, Int, ())
type PTGNode = (Int, NodeTag)
type PTG = Gr NodeTag ()
type Worklist = Seq Instruction
-- | Define a dependency graph.  The keys are the unique IDs of
-- locations.  The value for each ID is the set of instructions that
-- need to be re-processed when a new edge is added to the points-to graph
-- with that node ID as its source.
type DepGraph = IntMap (Set Instruction)

data Andersen = Andersen PTG

instance PointsToAnalysis Andersen where
  mayAlias (Andersen _) _ _ = True
  pointsTo = andersenPointsTo

andersenPointsTo :: (IsValue a) => Andersen -> a -> Set Value
andersenPointsTo (Andersen g) v =
  S.fromList $ map (unloc . lab g) (suc g (valueUniqueId v))
  where
    unloc (Just (Location l)) = l


-- | Run the points-to analysis and return an object that is an
-- instance of PointsToAnalysis, which can be used to query the
-- results.
runPointsToAnalysis :: Module -> Andersen
runPointsToAnalysis m = Andersen g
  where
    fs = moduleDefinedFunctions m
    blocks = concatMap functionBody fs
    insts = concatMap basicBlockInstructions blocks
    (globalLocations, globalEdges) = getGlobalLocations m
    argumentLocations = foldr extractArgs [] fs
    (localLocations, edgeInducers) = foldr extractLocations ([], []) insts
    allLocations = concat [globalLocations, argumentLocations, localLocations]
    -- The initial graph contains all locations in the program, along
    -- with edges induced by global initializers.
    graph0 = mkGraph allLocations globalEdges
    worklist0 = Seq.fromList edgeInducers
    g = saturate IM.empty worklist0 graph0

extractLocations :: Instruction
                    -> ([PTGNode], [Instruction])
                    -> ([PTGNode], [Instruction])
extractLocations i acc@(ptgNodes, insts) = case i of
  AllocaInst { instructionUniqueId = uid } ->
    ((uid, Location (Value i)) : ptgNodes, insts)
  StoreInst {} -> (ptgNodes, i : insts)
  CallInst {} -> (ptgNodes, i : insts)
  InvokeInst {} -> (ptgNodes, i : insts)
  _ -> acc

-- | Saturate the points-to graph using a worklist algorithm.
saturate :: DepGraph -> Worklist -> PTG -> PTG
saturate dg worklist g = case viewl worklist of
  EmptyL -> g
  itm :< rest -> case itm of
    StoreInst { storeValue = val, storeAddress = dest } ->
      case isPointerType (valueType val) of
        False -> saturate dg rest g
        True -> addStoreEdges dg itm val dest rest g
    CallInst { callFunction = cf, callArguments = args } ->
      addCallEdges dg rest g itm cf (map fst args)
    InvokeInst { invokeFunction = cf, invokeArguments = args } ->
      addCallEdges dg rest g itm cf (map fst args)
    _ -> error ("Unexpected instruction type in Andersen saturation: " ++ show itm)

-- | A call is essentially a copy of a pointer in the caller to an
-- argument node (which will later be copied in the callee).
--
-- Start by zipping together the actual arguments and formal argument
-- lists.  Filter out the non-pointer entries.
--
-- FIXME: Handle var-arg functions somehow
addCallEdges :: DepGraph -> Worklist -> PTG -> Instruction -> Value -> [Value] -> PTG
addCallEdges dg worklist g itm calledFunc args =
  case null newEdges of
    False -> saturate dg2 worklist' g'
    True -> saturate dg2 worklist g
  where
    (possibleCallees, dg1) = case valueContent calledFunc of
      -- Direct call
      FunctionC f -> ([ f ], dg)
      -- Indirect call: figure out what this value could possibly
      -- point to by consulting the points-to graph.  Note that this
      -- actually induces an extra dependency in the DepGraph (this is
      -- on-the-fly callgraph construction).
      _ -> undefined
    possibleFormalLists = map functionParameters possibleCallees
    formalsByPosition = transpose possibleFormalLists
    allActualFormalMap = zip args formalsByPosition
    pointerActualFormalMap = keepPointerParams allActualFormalMap
    keepPointerParams = filter (isPointerType . valueType . fst)
    (dg2, locMap) = mapAccumR getLocs dg1 pointerActualFormalMap
    -- The locations pointed to by the actuals are the targets of new
    -- points-to edges, while the locations of the formals are the
    -- sources of the edges.
    getLocs depGraph (actual, formals) =
      let (depGraph', actualLocs) = getLocationsReferencedBy g itm depGraph actual
          (depGraph'', formalLocs) = mapAccumR (getLocationsReferencedBy g itm) depGraph' formals
      in (depGraph'', zip actualLocs formalLocs)

    newEdges = locMapToEdges g locMap
    usedSrcs = foldl' accumUsedSrcs S.empty newEdges
    newWorklistItems = affectedInstructions usedSrcs dg2
    worklist' = worklist >< Seq.fromList newWorklistItems
    g' = foldl' (flip (&)) g newEdges


-- | The @locMap@ is an assoc list mapping the locations of actual
-- parameters to all of the potential formal parameter locations they
-- could correspond to.  The locations of actuals are the *targets* of
-- points-to graph edges, while the locations of formals are the
-- *sources*.
locMapToEdges g locMap = undefined
-- Fold over the loc-map to deal with each argument, then use an inner
-- fold over the sources and start identifying/checking edges.
-- Alternatively, use repeat to "copy" the targets to each pair of
-- sources and concat them. to use a single fold.


accumUsedSrcs :: Set Int -> Context a b -> Set Int
accumUsedSrcs acc (_, src, _, _) = S.insert src acc

-- | Handle adding edges induced by a @StoreInst@.
addStoreEdges :: DepGraph -> Instruction -> Value -> Value -> Worklist -> PTG -> PTG
addStoreEdges dg i val dest worklist g =
  -- Only update g and the worklist if there were new edges added.
  case null newEdges of
    False -> saturate dg2 worklist' g'
    -- IMPORTANT: Note that we need to propagate dependencies (the
    -- DepGraph) changes even if we didn't add any edges this time.
    -- Edges could be added later and we need to know all of the
    -- possible dependencies to come back.
    True -> saturate dg2 worklist g
  where
    (dg1, newTargets) = getLocationsReferencedBy g i dg val
    (dg2, newSrcs) = getLocationsReferencedBy g i dg1 dest
    newEdges = makeNewEdges g newSrcs newTargets
    usedSrcs = foldl' accumUsedSrcs S.empty newEdges
    newWorklistItems = affectedInstructions usedSrcs dg2
    worklist' = worklist >< Seq.fromList newWorklistItems
    g' = foldl' (flip (&)) g newEdges

affectedInstructions :: Set Int -> DepGraph -> [Instruction]
affectedInstructions usedSrcs dg = S.toList instSet
  where
    instSet = S.fold findAffected S.empty usedSrcs
    findAffected nodeId acc = S.union acc (IM.findWithDefault S.empty nodeId dg)

-- | Determine which edges need to be added to the graph, based on the
-- set of discovered sources and targets.  Only new edges are
-- returned.
makeNewEdges :: PTG -> [Node] -> [Node] -> [Context NodeTag ()]
makeNewEdges g newSrcs newTargets =
  mapMaybe toContext newSrcs
  where
    notInGraph src tgt = notElem tgt (suc g src)
    -- | Create a new context for each src in the graph.  A context is
    -- only created if it adds new targets.  This is to make worklist
    -- management easier; if there are no new edges, the worklist is
    -- not updated.
    --
    -- Note that all edges for each source are added at once to the
    -- context.  Using separate contexts for each new edge would
    -- require changes later, otherwise some contexts could overwrite
    -- others, losing edges.
    toContext src =
      let targets = filter (notInGraph src) newTargets
          (adjIn, n, lbl, adjOut) = context g src
      in case targets of
        [] -> Nothing
        _ -> Just (adjIn, n, lbl, map (\t->((), t)) targets ++ adjOut)

-- | Given a @Value@ that is an operand of a @StoreInst@, find all of
-- the locations in the points-to graph that it refers to.  This means
-- that the function starts at the given value @v@ and _dereferences_
-- each @LoadInst@ in the @Value@.
--
-- As it progresses, the function also updates the DepGraph that
-- allows for reverse lookups (which instructions depend on the edges
-- coming from a given node).  This DepGraph is used to determine
-- which instructions need to be added to the Worklist due to
-- additional edges added to the points-to graph each iteration.  Note
-- that the DepGraph should be updated *before* it is used to compute
-- the list of affected instructions.
--
-- This function treats BitcastInsts as no-ops (i.e., it assumes the
-- types work out and only deals with memory references).
--
-- Algorithm: walk the chain of loads/geps/bitcasts to find the
-- location being loaded from.  Keep track of the number of
-- dereferences (loads).  Take that many steps (across all outgoing
-- edges) from the location in the points-to graph.
getLocationsReferencedBy :: IsValue a
                            => PTG
                            -> Instruction
                            -> DepGraph
                            -> a
                            -> (DepGraph, [Node])
getLocationsReferencedBy g storeInst dg0 v = getLocs v 0
  where
    getLocs :: IsValue a => a -> Int -> (DepGraph, [Node])
    getLocs val derefCount = case valueContent val of
      -- This also needs to handle GEP instructions later, also
      -- (maybe) select, extractelement, and extractvalue.  This also
      -- needs to be careful around bitcasts of non-pointer types
      -- (integers turned into pointers, for example)
      InstructionC (LoadInst { loadAddress = addr }) ->
        getLocs addr (derefCount + 1)
      InstructionC (BitcastInst { castedValue = cv }) ->
        getLocs cv derefCount
      -- In this fallback case, @val@ should be a node in the
      -- Points-to graph.  Collect everything @derefCount@ steps from
      -- it and return that, along with the updated DepGraph.  The
      -- depgraph needs to be updated with an entry for each non-leaf
      -- node.
      _ -> collectLocationNodes g dg0 storeInst val derefCount

collectLocationNodes :: IsValue a
                        => PTG
                        -> DepGraph
                        -> Instruction
                        -> a
                        -> Int
                        -> (DepGraph, [Node])
collectLocationNodes g dg0 storeInst val derefCount =
  collect dg0 derefCount [valueUniqueId val]
  where
    addDep dg n = IM.insertWith S.union n (S.singleton storeInst) dg
    collect dg 0 resultNodes = (dg, resultNodes)
    collect dg remainingHops startNodes =
      let nextNodes = concatMap (suc g) startNodes
          dg' = foldl' addDep dg startNodes
      in collect dg' (remainingHops - 1) nextNodes


-- | This only re-allocates the small part of the list each iteration,
-- so should remain efficient.
{'fraction_non_alphanumeric': 0.081463173504696, 'fraction_numerical': 0.0020761245674740486, 'mean_word_length': 2.8222180238050254, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 129, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5904114', 'n_tokens_mistral': 5899, 'n_tokens_neox': 5758, 'n_words': 2400}