source/FrameSet.h

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#pragma once

#include <initializer_list>
#include <ostream>

#include <json/json.h>

#include <AbstractDomain.h>
#include <PatriciaTreeMapAbstractPartition.h>

#include <mariana-trench/FlattenIterator.h>
#include <mariana-trench/Frame.h>
#include <mariana-trench/GroupHashedSetAbstractDomain.h>

namespace marianatrench {

/* Represents a set of frames with the same kind. */
class FrameSet final : public sparta::AbstractDomain<FrameSet> {
 private:
  using Set =
      GroupHashedSetAbstractDomain<Frame, Frame::GroupHash, Frame::GroupEqual>;
  using CallPositionToSetMap =
      sparta::PatriciaTreeMapAbstractPartition<const Position*, Set>;
  using CalleeToCallPositionToSetMap = sparta::
      PatriciaTreeMapAbstractPartition<const Method*, CallPositionToSetMap>;

 private:
  // Iterator based on `FlattenIterator`.

  struct CallPositionToSetMapDereference {
    static Set::iterator begin(const std::pair<const Position*, Set>& pair) {
      return pair.second.begin();
    }
    static Set::iterator end(const std::pair<const Position*, Set>& pair) {
      return pair.second.end();
    }
  };

  using CallPositionToSetMapIterator = FlattenIterator<
      /* OuterIterator */ CallPositionToSetMap::MapType::iterator,
      /* InnerIterator */ Set::iterator,
      CallPositionToSetMapDereference>;

  struct CalleeToCallPositionToSetMapDereference {
    static CallPositionToSetMapIterator begin(
        const std::pair<const Method*, CallPositionToSetMap>& pair) {
      return CallPositionToSetMapIterator(
          pair.second.bindings().begin(), pair.second.bindings().end());
    }
    static CallPositionToSetMapIterator end(
        const std::pair<const Method*, CallPositionToSetMap>& pair) {
      return CallPositionToSetMapIterator(
          pair.second.bindings().end(), pair.second.bindings().end());
    }
  };

  using ConstIterator = FlattenIterator<
      /* OuterReference */ CalleeToCallPositionToSetMap::MapType::iterator,
      /* InnerIterator */ CallPositionToSetMapIterator,
      CalleeToCallPositionToSetMapDereference>;

 public:
  // C++ container concept member types
  using iterator = ConstIterator;
  using const_iterator = ConstIterator;
  using value_type = Frame;
  using difference_type = std::ptrdiff_t;
  using size_type = std::size_t;
  using const_reference = const Frame&;
  using const_pointer = const Frame*;

 private:
  explicit FrameSet(
      const Kind* MT_NULLABLE kind,
      CalleeToCallPositionToSetMap map)
      : kind_(kind), map_(std::move(map)) {}

 public:
  /* Create the bottom (i.e, empty) frame set. */
  FrameSet() : kind_(nullptr), map_(CalleeToCallPositionToSetMap::bottom()) {}

  explicit FrameSet(std::initializer_list<Frame> frames);

  FrameSet(const FrameSet&) = default;
  FrameSet(FrameSet&&) = default;
  FrameSet& operator=(const FrameSet&) = default;
  FrameSet& operator=(FrameSet&&) = default;

  static FrameSet bottom() {
    return FrameSet(/* kind */ nullptr, CalleeToCallPositionToSetMap::bottom());
  }

  static FrameSet top() {
    return FrameSet(/* kind */ nullptr, CalleeToCallPositionToSetMap::top());
  }

  bool is_bottom() const override {
    return map_.is_bottom();
  }

  bool is_top() const override {
    return map_.is_top();
  }

  void set_to_bottom() override {
    kind_ = nullptr;
    map_.set_to_bottom();
  }

  void set_to_top() override {
    kind_ = nullptr;
    map_.set_to_top();
  }

  bool empty() const {
    return map_.is_bottom();
  }

  const Kind* MT_NULLABLE kind() const {
    return kind_;
  }

  bool is_artificial_sources() const {
    return kind_ == Kinds::artificial_source();
  }

  void add(const Frame& frame);

  bool leq(const FrameSet& other) const override;

  bool equals(const FrameSet& other) const override;

  void join_with(const FrameSet& other) override;

  void widen_with(const FrameSet& other) override;

  void meet_with(const FrameSet& other) override;

  void narrow_with(const FrameSet& other) override;

  void difference_with(const FrameSet& other);

  void map(const std::function<void(Frame&)>& f);

  void filter(const std::function<bool(const Frame&)>& predicate);

  ConstIterator begin() const {
    return ConstIterator(map_.bindings().begin(), map_.bindings().end());
  }

  ConstIterator end() const {
    return ConstIterator(map_.bindings().end(), map_.bindings().end());
  }

  void add_inferred_features(const FeatureMayAlwaysSet& features);

  LocalPositionSet local_positions() const;

  void add_local_position(const Position* position);

  void set_local_positions(const LocalPositionSet& positions);

  void add_inferred_features_and_local_position(
      const FeatureMayAlwaysSet& features,
      const Position* MT_NULLABLE position);

  /**
   * Propagate the taint from the callee to the caller.
   *
   * Return bottom if the taint should not be propagated.
   */
  FrameSet propagate(
      const Method* caller,
      const Method* callee,
      const AccessPath& callee_port,
      const Position* call_position,
      int maximum_source_sink_distance,
      Context& context,
      const std::vector<const DexType * MT_NULLABLE>& source_register_types,
      const std::vector<std::optional<std::string>>& source_constant_arguments)
      const;

  /* Return the set of leaf frames with the given position. */
  FrameSet attach_position(const Position* position) const;

  /**
   * Convert the kind of these frames into the given kind.
   *
   * Return A new FrameSet, identical in everyway except in "kind".
   */
  FrameSet with_kind(const Kind* kind) const;

  template <class T>
  std::unordered_map<T, std::vector<std::reference_wrapper<const Frame>>>
  partition_map(const std::function<T(const Frame&)>& map) const;

  static FrameSet from_json(const Json::Value& value, Context& context);
  Json::Value to_json() const;

  friend std::ostream& operator<<(std::ostream& out, const FrameSet& frames);

 private:
  Frame propagate_frames(
      const Method* callee,
      const AccessPath& callee_port,
      const Position* call_position,
      int maximum_source_sink_distance,
      Context& context,
      const std::vector<const DexType * MT_NULLABLE>& source_register_types,
      const std::vector<std::optional<std::string>>& source_constant_arguments,
      std::vector<std::reference_wrapper<const Frame>> frames,
      std::vector<const Feature*>& via_type_of_features_added) const;

  FrameSet propagate_crtex_frames(
      const Method* callee,
      const AccessPath& callee_port,
      const Position* call_position,
      int maximum_source_sink_distance,
      Context& context,
      const std::vector<const DexType * MT_NULLABLE>& source_register_types,
      std::vector<std::reference_wrapper<const Frame>> frames) const;

 private:
  const Kind* MT_NULLABLE kind_;
  CalleeToCallPositionToSetMap map_;
};

} // namespace marianatrench