include/FiniteAbstractDomain.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 <bitset>
#include <cstddef>
#include <functional>
#include <initializer_list>
#include <ostream>
#include <sstream>
#include <type_traits>
#include <unordered_map>
#include <utility>

#include "AbstractDomain.h"

namespace sparta {

/*
 * This is the general interface for arbitrary encodings of a lattice. 'Element'
 * is the type of the symbolic names for the lattice elements and 'Encoding' is
 * the type of the actual encoding.
 */
template <typename Element, typename Encoding>
class LatticeEncoding {
 public:
  virtual ~LatticeEncoding() {}

  virtual Encoding encode(const Element& element) const = 0;

  virtual Element decode(const Encoding& encoding) const = 0;

  virtual bool is_bottom(const Encoding& element) const = 0;

  virtual bool is_top(const Encoding& element) const = 0;

  virtual bool equals(const Encoding& x, const Encoding& y) const = 0;

  virtual bool leq(const Encoding& x, const Encoding& y) const = 0;

  virtual Encoding join(const Encoding& x, const Encoding& y) const = 0;

  virtual Encoding meet(const Encoding& x, const Encoding& y) const = 0;

  virtual Encoding bottom() const = 0;

  virtual Encoding top() const = 0;
};

/*
 * Example usage:
 *
 *   Encoding the following lattice using bit vectors:
 *
 *              TOP
 *             /   \
 *            A     B
 *             \   /
 *             BOTTOM
 *
 *   enum Elements {BOTTOM, A, B, TOP};
 *   using Lattice = BitVectorLattice<Elements, 4, std::hash<int>>;
 *   Lattice lattice({BOTTOM, A, B, TOP},
 *                   {{BOTTOM, A}, {BOTTOM, B}, {A, TOP}, {B, TOP}});
 *   using Domain =
 *       FiniteAbstractDomain<Elements, Lattice, Lattice::Encoding, &lattice>;
 *   ...
 *   Domain a(A), b(B);
 *   Domain x = a.join(b);
 *   ...
 *
 * Note: since 'lattice' is a template argument, this object must be statically
 * defined, for example as a global variable. The lattice is instantiated just
 * once at startup time.
 *
 */
template <typename Element,
          typename Lattice,
          typename Encoding,
          Lattice* lattice>
class FiniteAbstractDomain final
    : public AbstractDomain<
          FiniteAbstractDomain<Element, Lattice, Encoding, lattice>> {
 public:
  ~FiniteAbstractDomain() {
    // The destructor is the only method that is guaranteed to be created when a
    // class template is instantiated. This is a good place to perform all the
    // sanity checks on the template parameters.
    static_assert(
        std::is_base_of<LatticeEncoding<Element, Encoding>, Lattice>::value,
        "Lattice doesn't derive from LatticeEncoding");
  }

  /*
   * A default constructor is required in the AbstractDomain specification.
   */
  FiniteAbstractDomain() : m_encoding(lattice->top()) {}

  explicit FiniteAbstractDomain(const Element& element)
      : m_encoding(lattice->encode(element)) {}

  Element element() const { return lattice->decode(m_encoding); }

  bool is_bottom() const override { return lattice->is_bottom(m_encoding); }

  bool is_top() const override { return lattice->is_top(m_encoding); }

  bool leq(const FiniteAbstractDomain& other) const override {
    return lattice->leq(m_encoding, other.m_encoding);
  }

  bool equals(const FiniteAbstractDomain& other) const override {
    return lattice->equals(m_encoding, other.m_encoding);
  }

  void set_to_bottom() override { m_encoding = lattice->bottom(); }

  void set_to_top() override { m_encoding = lattice->top(); }

  void join_with(const FiniteAbstractDomain& other) override {
    m_encoding = lattice->join(m_encoding, other.m_encoding);
  }

  void widen_with(const FiniteAbstractDomain& other) override {
    join_with(other);
  }

  void meet_with(const FiniteAbstractDomain& other) override {
    m_encoding = lattice->meet(m_encoding, other.m_encoding);
  }

  void narrow_with(const FiniteAbstractDomain& other) override {
    meet_with(other);
  }

  static FiniteAbstractDomain bottom() {
    return FiniteAbstractDomain(lattice->bottom());
  }

  static FiniteAbstractDomain top() {
    return FiniteAbstractDomain(lattice->top());
  }

 private:
  FiniteAbstractDomain(const Encoding& encoding) : m_encoding(encoding) {}

  Encoding m_encoding;
};

} // namespace sparta

template <typename Element,
          typename Lattice,
          typename Encoding,
          Lattice* lattice>
inline std::ostream& operator<<(
    std::ostream& o,
    const typename sparta::
        FiniteAbstractDomain<Element, Lattice, Encoding, lattice>& x) {
  o << x.element();
  return o;
}

namespace sparta {

namespace fad_impl {

/*
 * Our encoding of lattices is based on the following paper that proposes an
 * efficient representation based on bit vectors:
 *
 *   H. Aït-Kaci, R. Boyer, P. Lincoln, R. Nasr. Efficient implementation of
 *   lattice operations. In ACM Transactions on Programming Languages and
 *   Systems (TOPLAS), Volume 11, Issue 1, Jan. 1989, pages 115-146.
 *
 * The approach described in the paper only works with the Meet operation. The
 * idea is to represent the Hasse diagram of a lattice using a Boolean matrix,
 * as shown below:
 *
 *         d                          a  b  c  d
 *        / \                      a  0  0  0  0
 *       b   c                     b  1  0  0  0
 *        \ /                      c  1  0  0  0
 *         a                       d  0  1  1  0
 *
 * This matrix represents the "immediately greater than" relation in the
 * lattice. The technique consists of computing the reflexive and transitive
 * closure of that relation. Then, an element can be encoded by its
 * corresponding row (i.e., a bit vector) in the resulting matrix. Computing the
 * Meet simply amounts to performing the bitwise And operation on the bit
 * vector representation. For the example above that gives:
 *
 * Reflexive-transitive closure:
 *
 *               a  b  c  d                b Meet c = 1100 & 1010
 *            a  1  0  0  0                         = 1000
 *            b  1  1  0  0                         = a
 *            c  1  0  1  0
 *            d  1  1  1  1
 *
 * In order to compute the Join, we apply the same technique to the opposite
 * lattice, i.e., the lattice in which the order relation has been reversed and
 * the Top and Bottom elements have been swapped. The opposite lattice and the
 * corresponding Boolean matrix are constructed as follows:
 *
 *         a                          a  b  c  d
 *        / \                      a  0  1  1  0
 *       b   c                     b  0  0  0  1
 *        \ /                      c  0  0  0  1
 *         d                       d  0  0  0  0
 *
 * It can be easily seen that the Meet in the opposite lattice is exactly the
 * Join in the original lattice.
 *
 * Reflexive-transitive closure:
 *
 *               a  b  c  d                b Meet c = 0101 & 0011
 *            a  1  1  1  1                         = 0001
 *            b  0  1  0  1                         = d
 *            c  0  0  1  1                         = b Join c in the original
 *            d  0  0  0  1                           lattice
 *
 * The template parameter 'construct_opposite_lattice' specifies the lattice to
 * consider for the encoding.
 *
 * Note that constructing this representation has cubic time complexity in the
 * number of elements of the lattice. Since the construction is done only once
 * at startup time and finite lattices built this way are usually small, this
 * should not be a problem in practice.
 *
 */
template <typename Element,
          size_t cardinality,
          bool construct_opposite_lattice,
          typename Hash,
          typename Equal>
class BitVectorSemiLattice final {
 public:
  // The size of a bitset structure is a compile-time constant, hence the need
  // for the 'cardinality' parameter.
  using Encoding = std::bitset<cardinality>;

  BitVectorSemiLattice() = delete;

  /*
   * In order to construct the bit vector representation, the user provides the
   * complete set of elements in the lattice (including the Top and Bottom
   * elements) as well as the Hasse diagram of the partial order relation.
   */
  BitVectorSemiLattice(
      std::initializer_list<Element> elements,
      std::initializer_list<std::pair<Element, Element>> hasse_diagram) {
    RUNTIME_CHECK(elements.size() == cardinality,
                  invalid_argument()
                      << argument_name("elements")
                      << operation_name("BitVectorSemiLattice()"));

    // We assign each element of the lattice an index, so that we can construct
    // the Boolean matrix.
    Element index_to_element[cardinality];
    std::unordered_map<Element, size_t, Hash, Equal> element_to_index;

    std::copy(elements.begin(), elements.end(), index_to_element);
    for (size_t i = 0; i < cardinality; ++i) {
      element_to_index[index_to_element[i]] = i;
    }

    // We populate the Boolean matrix by traversing the Hasse diagram of the
    // partial order.
    Encoding matrix[cardinality];
    for (auto pair : hasse_diagram) {
      // The Hasse diagram provided by the user describes the partial order in
      // the original lattice. We need to normalize the representation when the
      // opposite lattice is considered.
      if (construct_opposite_lattice) {
        std::swap(pair.first, pair.second);
      }

      // If y is immediately greater than x in the partial order considered,
      // then matrix[y][x] = 1.
      auto x_it = element_to_index.find(pair.first);
      auto y_it = element_to_index.find(pair.second);
      RUNTIME_CHECK(x_it != element_to_index.end() &&
                        y_it != element_to_index.end(),
                    internal_error());
      matrix[y_it->second][x_it->second] = true;
    }

    // We first compute the reflexive closure of the "immediately greater than"
    // relation in the lattice considered.
    for (size_t i = 0; i < cardinality; ++i) {
      matrix[i][i] = true;
    }

    // Then we compute the transitive closure of the "immediately greater than"
    // relation in the lattice considered, using Warshall's algorithm.
    for (size_t k = 0; k < cardinality; ++k) {
      for (size_t i = 0; i < cardinality; ++i) {
        for (size_t j = 0; j < cardinality; ++j) {
          matrix[i][j] = matrix[i][j] || (matrix[i][k] && matrix[k][j]);
        }
      }
    }

    // The last step is to assign a bit vector representation to each element in
    // the lattice considered, i.e. the corresponding row in the Boolean matrix.
    // We also maintain a reverse table for decoding purposes.
    for (size_t i = 0; i < cardinality; ++i) {
      Element element = index_to_element[i];
      Encoding encoding = matrix[i];
      m_element_to_encoding[element] = encoding;
      m_encoding_to_element[encoding] = element;
      // We identify the Bottom and Top elements on the fly.
      if (is_bottom(encoding)) {
        m_bottom = encoding;
      }
      if (is_top(encoding)) {
        m_top = encoding;
      }
    }

    // Make sure that we obtain a semi-lattice.
    sanity_check(&index_to_element[0]);
  }

  Encoding encode(const Element& element) const {
    auto it = m_element_to_encoding.find(element);
    RUNTIME_CHECK(it != m_element_to_encoding.end(), undefined_operation());
    return it->second;
  }

  Element decode(const Encoding& encoding) const {
    auto it = m_encoding_to_element.find(encoding);
    RUNTIME_CHECK(it != m_encoding_to_element.end(), undefined_operation());
    return it->second;
  }

  bool is_bottom(const Encoding& x) const {
    // In the lower semi-lattice representation the Bottom element is the unique
    // bit vector that has only one bit set to 1, whereas in the opposite
    // semi-lattice it has all its bits set to 1.
    return construct_opposite_lattice ? x.all() : (x.count() == 1);
  }

  bool is_top(const Encoding& x) const {
    // The Top element is defined as the the dual of Bottom.
    return construct_opposite_lattice ? (x.count() == 1) : x.all();
  }

  Encoding bottom() const { return m_bottom; }

  Encoding top() const { return m_top; }

 private:
  // This sanity check verifies that the bitwise And of any two pairs of
  // elements (i.e., the Meet or the Join of those elements depending on the
  // lattice considered) corresponds to an actual element in the lattice.
  // In other words, this procedure makes sure that the input Hasse diagram
  // defines a semi-lattice.
  void sanity_check(Element* index_to_element) {
    // We count the number of bit vectors that have all their bits set to one.
    size_t all_bits_are_set = 0;
    // We count the number of bit vectors that have only one bit set to one.
    size_t one_bit_is_set = 0;
    for (size_t i = 0; i < cardinality; ++i) {
      Encoding x = m_element_to_encoding[index_to_element[i]];
      if (x.all()) {
        ++all_bits_are_set;
      }
      if (x.count() == 1) {
        ++one_bit_is_set;
      }
      for (size_t j = 0; j < cardinality; ++j) {
        Encoding y = m_element_to_encoding[index_to_element[j]];
        RUNTIME_CHECK(m_encoding_to_element.find(x & y) !=
                          m_encoding_to_element.end(),
                      internal_error());
      }
    }
    RUNTIME_CHECK(all_bits_are_set == 1 && one_bit_is_set == 1,
                  internal_error()
                      << error_msg("Missing or duplicate extremal element"));
  }

  std::unordered_map<Element, Encoding, Hash, Equal> m_element_to_encoding;
  std::unordered_map<Encoding, Element> m_encoding_to_element;
  Encoding m_bottom;
  Encoding m_top;
};

} // namespace fad_impl

/*
 * A lattice maintains two semi-lattices internally, always use opposite
 * semi-lattice representation and calculate corresponding lower semi-lattice
 * when needed
 */
template <typename Element,
          size_t cardinality,
          typename Hash = std::hash<Element>,
          typename Equal = std::equal_to<Element>>
class BitVectorLattice final
    : public LatticeEncoding<Element, std::bitset<cardinality>> {
 public:
  using Encoding = std::bitset<cardinality>;

  ~BitVectorLattice() {
    // The destructor is the only method that is guaranteed to be created when a
    // class template is instantiated. This is a good place to perform all the
    // sanity checks on the template parameters.
    static_assert(std::is_default_constructible<Element>::value,
                  "Element is not default constructible");
    static_assert(std::is_copy_constructible<Element>::value,
                  "Element is not copy constructible");
    static_assert(std::is_copy_assignable<Element>::value,
                  "Element is not copy assignable");
  }

  BitVectorLattice() = delete;

  BitVectorLattice(
      std::initializer_list<Element> elements,
      std::initializer_list<std::pair<Element, Element>> hasse_diagram)
      : m_lower_semi_lattice(elements, hasse_diagram),
        m_opposite_semi_lattice(elements, hasse_diagram) {}

  Encoding encode(const Element& element) const override {
    // In a standard fixpoint computation the Join is by far the dominant
    // operation. Hence, we favor the opposite semi-lattice encoding whenever we
    // construct a domain element.
    return m_opposite_semi_lattice.encode(element);
  }

  // Default use opposite semi-lattice for decoding.
  Element decode(const Encoding& encoding) const override {
    return m_opposite_semi_lattice.decode(encoding);
  }

  Element decode_lower(const Encoding& encoding) const {
    return m_lower_semi_lattice.decode(encoding);
  }

  bool is_bottom(const Encoding& x) const override { return x.all(); }

  bool is_top(const Encoding& x) const override { return x.count() == 1; }

  bool equals(const Encoding& x, const Encoding& y) const override {
    return x == y;
  }

  bool leq(const Encoding& x, const Encoding& y) const override {
    return (x & y) == y;
  }

  Encoding join(const Encoding& x, const Encoding& y) const override {
    return x & y;
  }

  Encoding meet(const Encoding& x, const Encoding& y) const override {
    // In order to perform the Meet, we need to calculate corresponding lower
    // semi-lattice encoding, and switch back to opposite semi-lattice encoding
    // before returning.
    auto x_lower = get_lower_encoding(x);
    auto y_lower = get_lower_encoding(y);
    Encoding lower_encoding = x_lower & y_lower;
    return get_opposite_encoding(lower_encoding);
  }

  Encoding bottom() const override { return m_opposite_semi_lattice.bottom(); }

  Encoding top() const override { return m_opposite_semi_lattice.top(); }

 private:
  Encoding get_lower_encoding(const Encoding& x) const {
    const Element& element = decode(x);
    return m_lower_semi_lattice.encode(element);
  }

  Encoding get_opposite_encoding(const Encoding& x) const {
    const Element& element = decode_lower(x);
    return m_opposite_semi_lattice.encode(element);
  }

  fad_impl::BitVectorSemiLattice<Element,
                                 cardinality,
                                 /* construct_opposite_lattice */ false,
                                 Hash,
                                 Equal>
      m_lower_semi_lattice;
  fad_impl::BitVectorSemiLattice<Element,
                                 cardinality,
                                 /* construct_opposite_lattice */ true,
                                 Hash,
                                 Equal>
      m_opposite_semi_lattice;
};

} // namespace sparta