text: # Copyright (c) 2016-present, Facebook, Inc.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  from collections import Counter
  from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
  
  from .models import SharedTextKind, TraceFrame, TraceFrameAnnotation, TraceKind
  from .trace_graph import TraceGraph
  
  
  class TrimmedTraceGraph(TraceGraph):
      """Represents a trimmed graph that is constructed from a bigger TraceGraph
      based on issues that have traces involving a set of affected files or
      directories.
      """
  
      def __init__(
          self, affected_files: List[str], affected_issues_only: bool = False
      ) -> None:
          """Creates an empty TrimmedTraceGraph.
          """
          super().__init__()
          self._affected_files = affected_files
          self._affected_issues_only = affected_issues_only
          self._visited_trace_frame_ids: Set[int] = set()
  
      def populate_from_trace_graph(self, graph: TraceGraph) -> None:
          """Populates this graph from the given one based on affected_files
          """
          # Track which trace frames have been visited as we populate the full
          # traces of the graph.
          self._visited_trace_frame_ids: Set[int] = set()
  
          self._populate_affected_issues(graph)
  
          if not self._affected_issues_only:
              # Finds issues from the conditions and saves them.
              # Also saves traces that have been trimmed to the affected
              # conditions.
              self._populate_issues_from_affected_trace_frames(graph)
  
              # Traces populated above may be missing all traces because
              # _populate_issues_from_affected_trace_frames only populates
              # traces that reach the affected conditions in one direction. We
              # may need to populate traces in other directions too.
              #
              # For example:
              #
              # Issue_x reaches affected_file_x via postcondition_x (forward
              # trace, i.e. trace leading to source). None of its backward
              # traces (leading to sinks) reach the affected files.
              #
              # _populate_issues_from_affected_trace_frames would have copied its
              # forward traces and trimmed it to those reaching postcondition_x.
              # We cannot blindly populate all forward traces in this case as
              # branches not leading to postcondition_x are unnecessary.
              #
              # However, in this specific example, all backward traces are needed
              # to give a complete picture of which sinks the issue reaches.
              # The following ensures that.
              for instance_id in self._issue_instances.keys():
                  first_hop_ids = self._issue_instance_trace_frame_assoc[instance_id]
                  fwd_trace_ids = {
                      tf_id
                      for tf_id in first_hop_ids
                      if self._trace_frames[tf_id].kind == TraceKind.POSTCONDITION
                  }
                  bwd_trace_ids = {
                      tf_id
                      for tf_id in first_hop_ids
                      if self._trace_frames[tf_id].kind == TraceKind.PRECONDITION
                  }
  
                  if len(fwd_trace_ids) == 0:
                      self._populate_issue_trace(
                          graph,
                          instance_id,
                          # pyre-fixme[6]: Expected `Optional[TraceKind]` for 3rd
                          #  param but got `(cls: TraceKind) -> Any`.
                          TraceKind.POSTCONDITION,
                      )
  
                  if len(bwd_trace_ids) == 0:
                      self._populate_issue_trace(
                          graph,
                          instance_id,
                          # pyre-fixme[6]: Expected `Optional[TraceKind]` for 3rd
                          #  param but got `(cls: TraceKind) -> Any`.
                          TraceKind.PRECONDITION,
                      )
  
          self._recompute_instance_properties()
  
      def _recompute_instance_properties(self):
          """Some properties of issue instances will be affected after trimming
          such as min trace length to leaves. This should be called after the
          trimming to re-compute these values.
          """
          callables_histo = Counter(
              inst.callable_id.local_id for inst in self._issue_instances.values()
          )
  
          for inst in self._issue_instances.values():
              inst.min_trace_length_to_sources = self._get_min_depth_to_sources(
                  inst.id.local_id
              )
              inst.min_trace_length_to_sinks = self._get_min_depth_to_sinks(
                  inst.id.local_id
              )
              inst.callable_count = callables_histo[inst.callable_id.local_id]
  
      def _get_min_depth_to_sources(self, instance_id: int) -> int:
          """Returns shortest depth to source from the issue instance. Instances
          have a pre-computed min_trace_length_to_source, but this can change
          after traces get trimmed from the graph. This re-computes it and
          returns the min.
          """
          source_ids = {
              source.id.local_id
              for source in (
                  # pyre-fixme[6]: Expected `SharedTextKind` for 2nd param but got
                  #  `(cls: SharedTextKind) -> Any`.
                  self.get_issue_instance_shared_texts(instance_id, SharedTextKind.SOURCE)
              )
          }
          first_hop_tf_ids = {
              tf_id
              for tf_id in self._issue_instance_trace_frame_assoc[instance_id]
              if self.get_trace_frame_from_id(tf_id).kind == TraceKind.POSTCONDITION
          }
          return self._get_min_depth(first_hop_tf_ids, source_ids)
  
      def _get_min_depth_to_sinks(self, instance_id: int) -> int:
          """See get_min_depths_to_sources."""
          sink_ids = {
              sink.id.local_id
              for sink in (
                  # pyre-fixme[6]: Expected `SharedTextKind` for 2nd param but got
                  #  `(cls: SharedTextKind) -> Any`.
                  self.get_issue_instance_shared_texts(instance_id, SharedTextKind.SINK)
              )
          }
          first_hop_tf_ids = {
              tf_id
              for tf_id in self._issue_instance_trace_frame_assoc[instance_id]
              if self.get_trace_frame_from_id(tf_id).kind == TraceKind.PRECONDITION
          }
          return self._get_min_depth(first_hop_tf_ids, sink_ids)
  
      def _get_min_depth(self, first_hop_tf_ids: Set[int], leaf_ids: Set[int]) -> int:
          min_depth = None
          for tf_id in first_hop_tf_ids:
              leaf_depths = self._trace_frame_leaf_assoc[tf_id]
              for (leaf_id, depth) in leaf_depths:
                  if leaf_id in leaf_ids and (min_depth is None or depth < min_depth):
                      min_depth = depth
          if min_depth is not None:
              return min_depth
          return 0
  
      def _populate_affected_issues(self, graph: TraceGraph) -> None:
          """Populates the trimmed graph with issues whose locations are in
          affected_files based on data in the input graph. Since these issues
          exist in the affected files, all traces are copied as well.
          """
          affected_instance_ids = [
              instance.id.local_id
              for instance in graph._issue_instances.values()
              if self._is_filename_prefixed_with(
                  graph.get_text(instance.filename_id), self._affected_files
              )
          ]
  
          for instance_id in affected_instance_ids:
              if instance_id in self._issue_instances:
                  continue
              self._populate_issue_and_traces(graph, instance_id)
  
      def _get_sink_names(self, graph: TraceGraph, instance_id: int) -> Set[str]:
          # pyre-fixme[9]: kind has type `SharedTextKind`; used as `(cls:
          #  SharedTextKind) -> Any`.
          kind: SharedTextKind = SharedTextKind.SINK
          sinks = graph.get_issue_instance_shared_texts(instance_id, kind)
          return {sink.contents for sink in sinks}
  
      def _get_source_names(self, graph: TraceGraph, instance_id: int) -> Set[str]:
          # pyre-fixme[9]: kind has type `SharedTextKind`; used as `(cls:
          #  SharedTextKind) -> Any`.
          kind: SharedTextKind = SharedTextKind.SOURCE
          sources = graph.get_issue_instance_shared_texts(instance_id, kind)
          return {source.contents for source in sources}
  
      def _get_leaf_names(self, graph: TraceGraph, instance_id: int) -> Set[str]:
          return self._get_source_names(graph, instance_id).union(
              self._get_sink_names(graph, instance_id)
          )
  
      def _get_leaf_names_from_pairs(
          self, graph: TraceGraph, leaf_pairs: Set[Tuple[int, Any]]
      ) -> Set[str]:
          return {
              leaf.contents
              for leaf in (graph._shared_texts[leaf_id] for (leaf_id, _) in leaf_pairs)
          }
  
      def _populate_issues_from_affected_trace_frames(self, graph: TraceGraph) -> None:
          """TraceFrames found in affected_files should be reachable via some
          issue instance. Follow traces of graph to find them and
          populate this TrimmedGraph with it.
          """
  
          initial_trace_frames = [
              trace_frame
              for trace_frame in graph._trace_frames.values()
              if self._is_filename_prefixed_with(
                  graph.get_text(trace_frame.filename_id), self._affected_files
              )
          ]
  
          self._populate_issues_from_affected_conditions(
              initial_trace_frames,
              graph,
              lambda trace_frame_id: (
                  graph._trace_frame_issue_instance_assoc[trace_frame_id]
              ),
              lambda trace_frame: (
                  [
                      graph._trace_frames[trace_frame_id]
                      for trace_frame_id in graph._trace_frames_rev_map[
                          (trace_frame.caller_id.local_id, trace_frame.caller_port)
                      ]
                      if graph._trace_frames[trace_frame_id].kind == trace_frame.kind
                  ]
              ),
              lambda instance_id: (self._get_leaf_names(graph, instance_id)),
              lambda trace_frame_id: (
                  self._get_leaf_names_from_pairs(
                      graph, graph._trace_frame_leaf_assoc[trace_frame_id]
                  )
              ),
              lambda instance, trace_frame: (
                  self.add_issue_instance_trace_frame_assoc(instance, trace_frame)
              ),
              lambda trace_frame_id: (
                  self._add_trace_frame(graph, graph._trace_frames[trace_frame_id])
              ),
              lambda initial_trace_frame_ids: (
                  self._populate_trace(graph, initial_trace_frame_ids)
              ),
          )
  
      def _populate_issues_from_affected_conditions(
          self,
          initial_conditions,
          graph: TraceGraph,
          get_issue_instances_from_condition_id,
          get_condition_parent,
          get_instance_leaves,
          get_condition_leaves,
          add_instance_condition_assoc,
          add_condition,
          add_traces,
      ) -> None:
          """Helper for populating reachable issue instances from the initial
          pre/postconditions. Also populates conditions/traces reachable from
          these instances. Traces are populated only in the direction this is
          called from: i.e. if initial_conditions are preconditions, only the
          backward trace is populated.
  
          Params:
          initial_conditions: The initial collection of pre/postconditions to
          start searching for issues from.
  
          graph: The trace graph to search for issues. Nodes/edges in this graph
          will be copied over to the local state
  
          get_issue_instances_from_condition_id: Function that returns all
          issue instances associated with a given a pre/postcondition id.
  
          get_condition_parent: Function that returns the parent condition of
          a given pre/postcondition. Given a pre/postcondition, p, its parent p',
          is the pre/postcondition that calls it, i.e. p.caller = p'.callee
  
          get_instance_leaves: Function that returns a collection of leaf names
          associated with the given issue instance ID.
  
          get_condition_leaves: Function that returns a collection of leaf names
          associated with the given (pre/post)condition ID.
  
          add_instance_condition_assoc: Function that takes in the issue
          instance and condition and adds the assoc between them.
  
          add_condition: Function that adds a condition to the graph given its
          id. This must add all the condition's assocs with the leaves because
          we don't filter out any condition-leaf assocs.
  
          add_traces: Function that takes a list of initial conditions and
          adds all conditions reachable from these to the graph.
          """
          visited: Dict[int, Set[str]] = {}
          que = [
              (cond, get_condition_leaves(cond.id.local_id))
              for cond in initial_conditions
          ]
  
          # Note that parent conditions may not transitively lead to the leaves
          # that its descendents lead to due to special-cased leaf filtering at
          # analysis time. When visiting each condition, we need to track the
          # leaves that we are visiting it from and only visit parent traces that
          # share common leaves along the path.
          while len(que) > 0:
              condition, leaves = que.pop()
              cond_id = condition.id.local_id
  
              if cond_id in visited:
                  leaves = leaves - visited[cond_id]
                  if len(leaves) == 0:
                      continue
                  else:
                      visited[cond_id].update(leaves)
              else:
                  visited[cond_id] = leaves
  
              # Found instance(s) related to the current condition. Yay.
              # This instance may have been found before, but process it again
              # anyway because we need to add the assoc with this condition.
              for instance_id in get_issue_instances_from_condition_id(cond_id):
                  # Check if the leaves (sources/sinks) of the issue reach
                  # the same leaves as the ones relevant to this condition.
                  instance = graph._issue_instances[instance_id]
                  issue_leaves = set(get_instance_leaves(instance.id.local_id))
                  common_leaves = issue_leaves.intersection(leaves)
                  if len(common_leaves) > 0:
                      if instance_id not in self._issue_instances:
                          self._populate_issue(graph, instance_id)
                      add_instance_condition_assoc(instance, condition)
  
              # Conditions that call this may have originated from other issues,
              # keep searching for parent conditions leading to this one.
              for next_cond in get_condition_parent(condition):
                  next_leaves = get_condition_leaves(next_cond.id.local_id)
                  common_leaves = leaves.intersection(next_leaves)
                  if len(leaves) > 0:
                      que.append((next_cond, common_leaves))
  
          # Add traces leading out from initial_conditions, and all visited
          # conditions leading back towards the issues.
          initial_condition_ids = [
              condition.id.local_id for condition in initial_conditions
          ]
          add_traces(initial_condition_ids)
          for condition_id in visited:
              add_condition(condition_id)
  
      def _populate_issue_and_traces(self, graph: TraceGraph, instance_id: int) -> None:
          """ Copies an issue over from the given trace graph, including all its
          traces and assocs.
          """
          self._populate_issue(graph, instance_id)
          self._populate_issue_trace(graph, instance_id)
  
      def _populate_issue_trace(
          self, graph: TraceGraph, instance_id: int, kind: Optional[TraceKind] = None
      ) -> None:
          trace_frame_ids = list(graph._issue_instance_trace_frame_assoc[instance_id])
          instance = graph._issue_instances[instance_id]
          filtered_ids = []
          for trace_frame_id in trace_frame_ids:
              frame = graph._trace_frames[trace_frame_id]
              if kind is None or kind == frame.kind:
                  self.add_issue_instance_trace_frame_assoc(instance, frame)
                  filtered_ids.append(trace_frame_id)
          self._populate_trace(graph, filtered_ids)
  
      def _populate_issue(self, graph: TraceGraph, instance_id: int) -> None:
          """Adds an issue to the trace graph along with relevant information
          pertaining to the issue (e.g. instance, fix_info, sources/sinks)
          The issue is identified by its corresponding instance's ID in the input
          trace graph.
          """
          instance = graph._issue_instances[instance_id]
          issue = graph._issues[instance.issue_id.local_id]
          self._populate_shared_text(graph, instance.message_id)
          self._populate_shared_text(graph, instance.filename_id)
          self._populate_shared_text(graph, instance.callable_id)
  
          self.add_issue_instance(instance)
          self.add_issue(issue)
  
          if instance_id in graph._issue_instance_fix_info:
              issue_fix_info = graph._issue_instance_fix_info[instance_id]
              self.add_issue_instance_fix_info(instance, issue_fix_info)
  
          for shared_text_id in graph._issue_instance_shared_text_assoc[instance_id]:
              shared_text = graph._shared_texts[shared_text_id]
              if shared_text_id not in self._shared_texts:
                  self.add_shared_text(shared_text)
              self.add_issue_instance_shared_text_assoc(instance, shared_text)
  
      def _populate_trace(self, graph: TraceGraph, trace_frame_ids: List[int]) -> None:
          """ Populates (from the given trace graph) the forward and backward
          traces reachable from the given traces (including input trace frames).
          Make sure to respect trace kind in successors
          """
          while len(trace_frame_ids) > 0:
              trace_frame_id = trace_frame_ids.pop()
              if trace_frame_id in self._visited_trace_frame_ids:
                  continue
  
              trace_frame = graph._trace_frames[trace_frame_id]
              self._add_trace_frame(graph, trace_frame)
              self._visited_trace_frame_ids.add(trace_frame_id)
  
              trace_frame_ids.extend(
                  [
                      next_frame.id.local_id
                      for next_frame in graph.get_next_trace_frames(trace_frame)
                      if next_frame.id.local_id not in self._visited_trace_frame_ids
                  ]
              )
  
      def _add_trace_frame(self, graph: TraceGraph, trace_frame: TraceFrame) -> None:
          """ Copies the trace frame from 'graph' to this (self) graph.
          Also copies all the trace_frame-leaf assocs since we don't
          know which ones are needed until we know the issue that reaches it
          """
          trace_frame_id = trace_frame.id.local_id
          self.add_trace_frame(trace_frame)
  
          annotations = graph.get_condition_annotations(trace_frame_id)
          for annotation in annotations:
              self._add_trace_annotation(graph, annotation)
  
          self._populate_shared_text(graph, trace_frame.filename_id)
          self._populate_shared_text(graph, trace_frame.caller_id)
          self._populate_shared_text(graph, trace_frame.callee_id)
          for (leaf_id, depth) in graph._trace_frame_leaf_assoc[trace_frame_id]:
              leaf = graph._shared_texts[leaf_id]
              if leaf_id not in self._shared_texts:
                  self.add_shared_text(leaf)
              self.add_trace_frame_leaf_assoc(trace_frame, leaf, depth)
  
      @staticmethod
      def _is_filename_prefixed_with(filename: str, prefixes: Iterable[str]) -> bool:
          return any([filename.startswith(p) for p in prefixes])
  
      def _populate_shared_text(self, graph, id) -> None:
          text = graph._shared_texts[id.local_id]
          if text.id.local_id not in self._shared_texts:
              self.add_shared_text(text)
  
      def _add_trace_annotation(
          self, graph: TraceGraph, annotation: TraceFrameAnnotation
      ) -> None:
          """Copies the annotation from 'graph' to this (self) graph.
          Also copies children TraceFrames of the annotation (if any). The
          parent TraceFrame of the annotation is NOT copied.
          """
          self.add_trace_annotation(annotation)
          children = graph.get_annotation_trace_frames(annotation.id.local_id)
          child_ids = [child.id.local_id for child in children]
          for child in children:
              self.add_trace_frame_annotation_trace_frame_assoc(annotation, child)
          self._populate_trace(graph, child_ids)
  