[clang][dataflow] Add parameterized map lattice.

This patchs adds a `MapLattice` template for lifting a lattice to a keyed map. A
typical use is for modeling variables in a scope with a partcular lattice.

Differential Revision: https://reviews.llvm.org/D116369

3 files changed, 297 insertions, 0 deletions

diff --git a/clang/include/clang/Analysis/FlowSensitive/MapLattice.h b/clang/include/clang/Analysis/FlowSensitive/MapLattice.h
new file mode 100644
index 000000000000..ff403f68b7c5
--- /dev/null
+++ b/clang/include/clang/Analysis/FlowSensitive/MapLattice.h
@@ -0,0 +1,140 @@
+//===------------------------ MapLattice.h ----------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+//  This file defines a parameterized lattice that maps keys to individual
+//  lattice elements (of the parameter lattice type). A typical usage is lifting
+//  a particular lattice to all variables in a lexical scope.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_ANALYSIS_FLOWSENSITIVE__MAPLATTICE_H
+#define LLVM_CLANG_ANALYSIS_FLOWSENSITIVE__MAPLATTICE_H
+
+#include <ostream>
+#include <string>
+#include <utility>
+
+#include "DataflowAnalysis.h"
+#include "clang/AST/Decl.h"
+#include "clang/Analysis/FlowSensitive/DataflowLattice.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringRef.h"
+
+namespace clang {
+namespace dataflow {
+
+/// A lattice that maps keys to individual lattice elements. When instantiated
+/// with an `ElementLattice` that is a bounded semi-lattice, `MapLattice` is
+/// itself a bounded semi-lattice, so long as the user limits themselves to a
+/// finite number of keys. In that case, `top` is (implicitly), the map
+/// containing all valid keys mapped to `top` of `ElementLattice`.
+///
+/// Requirements on `ElementLattice`:
+/// * Provides standard declarations of a bounded semi-lattice.
+template <typename Key, typename ElementLattice> class MapLattice {
+  using Container = llvm::DenseMap<Key, ElementLattice>;
+  Container C;
+
+public:
+  using key_type = Key;
+  using mapped_type = ElementLattice;
+  using value_type = typename Container::value_type;
+  using iterator = typename Container::iterator;
+  using const_iterator = typename Container::const_iterator;
+
+  MapLattice() = default;
+
+  explicit MapLattice(Container C) { C = std::move(C); }
+
+  // The `bottom` element is the empty map.
+  static MapLattice bottom() { return MapLattice(); }
+
+  void insert(const std::pair<const key_type, mapped_type> &P) { C.insert(P); }
+
+  void insert(std::pair<const key_type, mapped_type> &&P) {
+    C.insert(std::move(P));
+  }
+
+  unsigned size() const { return C.size(); }
+  bool empty() const { return C.empty(); }
+
+  iterator begin() { return C.begin(); }
+  iterator end() { return C.end(); }
+  const_iterator begin() const { return C.begin(); }
+  const_iterator end() const { return C.end(); }
+
+  // Equality is direct equality of underlying map entries. One implication of
+  // this definition is that a map with (only) keys that map to bottom is not
+  // equal to the empty map.
+  friend bool operator==(const MapLattice &LHS, const MapLattice &RHS) {
+    return LHS.C == RHS.C;
+  }
+
+  friend bool operator!=(const MapLattice &LHS, const MapLattice &RHS) {
+    return !(LHS == RHS);
+  }
+
+  bool contains(const key_type &K) const { return C.find(K) != C.end(); }
+
+  iterator find(const key_type &K) { return C.find(K); }
+  const_iterator find(const key_type &K) const { return C.find(K); }
+
+  mapped_type &operator[](const key_type &K) { return C[K]; }
+
+  /// If an entry exists in one map but not the other, the missing entry is
+  /// treated as implicitly mapping to `bottom`. So, the joined map contains the
+  /// entry as it was in the source map.
+  LatticeJoinEffect join(const MapLattice &Other) {
+    LatticeJoinEffect Effect = LatticeJoinEffect::Unchanged;
+    for (const auto &O : Other.C) {
+      auto It = C.find(O.first);
+      if (It == C.end()) {
+        C.insert(O);
+        Effect = LatticeJoinEffect::Changed;
+      } else if (It->second.join(O.second) == LatticeJoinEffect::Changed)
+        Effect = LatticeJoinEffect::Changed;
+    }
+    return Effect;
+  }
+};
+
+/// Convenience alias that captures the common use of map lattices to model
+/// in-scope variables.
+template <typename ElementLattice>
+using VarMapLattice = MapLattice<const clang::VarDecl *, ElementLattice>;
+
+template <typename Key, typename ElementLattice>
+std::ostream &
+operator<<(std::ostream &Os,
+           const clang::dataflow::MapLattice<Key, ElementLattice> &M) {
+  std::string Separator = "";
+  Os << "{";
+  for (const auto &E : M) {
+    Os << std::exchange(Separator, ", ") << E.first << " => " << E.second;
+  }
+  Os << "}";
+  return Os;
+}
+
+template <typename ElementLattice>
+std::ostream &
+operator<<(std::ostream &Os,
+           const clang::dataflow::VarMapLattice<ElementLattice> &M) {
+  std::string Separator = "";
+  Os << "{";
+  for (const auto &E : M) {
+    Os << std::exchange(Separator, ", ") << E.first->getName().str() << " => "
+       << E.second;
+  }
+  Os << "}";
+  return Os;
+}
+} // namespace dataflow
+} // namespace clang
+
+#endif // LLVM_CLANG_ANALYSIS_FLOWSENSITIVE__MAPLATTICE_H
diff --git a/clang/unittests/Analysis/FlowSensitive/CMakeLists.txt b/clang/unittests/Analysis/FlowSensitive/CMakeLists.txt
index 90c7be6b9068..753cf486953e 100644
--- a/clang/unittests/Analysis/FlowSensitive/CMakeLists.txt
+++ b/clang/unittests/Analysis/FlowSensitive/CMakeLists.txt
@@ -4,6 +4,7 @@ set(LLVM_LINK_COMPONENTS
   )
 
 add_clang_unittest(ClangAnalysisFlowSensitiveTests
+  MapLatticeTest.cpp
   SingleVarConstantPropagationTest.cpp
   TestingSupport.cpp
   TestingSupportTest.cpp
diff --git a/clang/unittests/Analysis/FlowSensitive/MapLatticeTest.cpp b/clang/unittests/Analysis/FlowSensitive/MapLatticeTest.cpp
new file mode 100644
index 000000000000..d3436e8f9496
--- /dev/null
+++ b/clang/unittests/Analysis/FlowSensitive/MapLatticeTest.cpp
@@ -0,0 +1,156 @@
+#include "clang/Analysis/FlowSensitive/MapLattice.h"
+#include "clang/Analysis/FlowSensitive/DataflowLattice.h"
+#include "llvm/Support/Error.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include <ostream>
+
+using namespace clang;
+using namespace dataflow;
+
+namespace {
+// A simple lattice for basic tests.
+class BooleanLattice {
+public:
+  BooleanLattice() : Value(false) {}
+  explicit BooleanLattice(bool B) : Value(B) {}
+
+  static BooleanLattice bottom() { return BooleanLattice(false); }
+
+  static BooleanLattice top() { return BooleanLattice(true); }
+
+  LatticeJoinEffect join(BooleanLattice Other) {
+    auto Prev = Value;
+    Value = Value || Other.Value;
+    return Prev == Value ? LatticeJoinEffect::Unchanged
+                         : LatticeJoinEffect::Changed;
+  }
+
+  friend bool operator==(BooleanLattice LHS, BooleanLattice RHS) {
+    return LHS.Value == RHS.Value;
+  }
+
+  friend bool operator!=(BooleanLattice LHS, BooleanLattice RHS) {
+    return LHS.Value != RHS.Value;
+  }
+
+  friend std::ostream &operator<<(std::ostream &Os, const BooleanLattice &B) {
+    Os << B.Value;
+    return Os;
+  }
+
+  bool value() const { return Value; }
+
+private:
+  bool Value;
+};
+} // namespace
+
+static constexpr int Key1 = 0;
+static constexpr int Key2 = 1;
+
+namespace {
+using ::testing::Pair;
+using ::testing::UnorderedElementsAre;
+
+TEST(MapLatticeTest, InsertWorks) {
+  MapLattice<int, BooleanLattice> Lattice;
+  Lattice.insert({Key1, BooleanLattice(false)});
+  Lattice.insert({Key2, BooleanLattice(false)});
+
+  EXPECT_THAT(Lattice, UnorderedElementsAre(Pair(Key1, BooleanLattice(false)),
+                                            Pair(Key2, BooleanLattice(false))));
+}
+
+TEST(MapLatticeTest, ComparisonWorks) {
+  MapLattice<int, BooleanLattice> Lattice1;
+  Lattice1.insert({Key1, BooleanLattice(true)});
+  Lattice1.insert({Key2, BooleanLattice(false)});
+  MapLattice<int, BooleanLattice> Lattice2 = Lattice1;
+  EXPECT_EQ(Lattice1, Lattice2);
+
+  Lattice2.find(Key2)->second = BooleanLattice(true);
+  EXPECT_NE(Lattice1, Lattice2);
+}
+
+TEST(MapLatticeTest, JoinChange) {
+  MapLattice<int, BooleanLattice> Lattice1;
+  Lattice1.insert({Key1, BooleanLattice(false)});
+  Lattice1.insert({Key2, BooleanLattice(false)});
+
+  MapLattice<int, BooleanLattice> Lattice2;
+  Lattice2.insert({Key1, BooleanLattice(true)});
+  Lattice2.insert({Key2, BooleanLattice(true)});
+
+  ASSERT_THAT(Lattice1,
+              UnorderedElementsAre(Pair(Key1, BooleanLattice(false)),
+                                   Pair(Key2, BooleanLattice(false))));
+
+  ASSERT_EQ(Lattice1.join(Lattice2), LatticeJoinEffect::Changed);
+  EXPECT_THAT(Lattice1, UnorderedElementsAre(Pair(Key1, BooleanLattice(true)),
+                                             Pair(Key2, BooleanLattice(true))));
+}
+
+TEST(MapLatticeTest, JoinEqNoChange) {
+  MapLattice<int, BooleanLattice> Lattice;
+  Lattice.insert({Key1, BooleanLattice(false)});
+  Lattice.insert({Key2, BooleanLattice(false)});
+
+  ASSERT_EQ(Lattice.join(Lattice), LatticeJoinEffect::Unchanged);
+  EXPECT_THAT(Lattice, UnorderedElementsAre(Pair(Key1, BooleanLattice(false)),
+                                            Pair(Key2, BooleanLattice(false))));
+}
+
+TEST(MapLatticeTest, JoinLtNoChange) {
+  MapLattice<int, BooleanLattice> Lattice1;
+  Lattice1.insert({Key1, BooleanLattice(false)});
+  Lattice1.insert({Key2, BooleanLattice(false)});
+
+  MapLattice<int, BooleanLattice> Lattice2;
+  Lattice2.insert({Key1, BooleanLattice(true)});
+  Lattice2.insert({Key2, BooleanLattice(true)});
+
+  ASSERT_THAT(Lattice1,
+              UnorderedElementsAre(Pair(Key1, BooleanLattice(false)),
+                                   Pair(Key2, BooleanLattice(false))));
+
+  ASSERT_THAT(Lattice2, UnorderedElementsAre(Pair(Key1, BooleanLattice(true)),
+                                             Pair(Key2, BooleanLattice(true))));
+
+  ASSERT_EQ(Lattice2.join(Lattice1), LatticeJoinEffect::Unchanged);
+  EXPECT_THAT(Lattice2, UnorderedElementsAre(Pair(Key1, BooleanLattice(true)),
+                                             Pair(Key2, BooleanLattice(true))));
+}
+
+TEST(MapLatticeTest, JoinDifferentDomainsProducesUnion) {
+  MapLattice<int, BooleanLattice> Lattice1;
+  Lattice1.insert({Key1, BooleanLattice(true)});
+  MapLattice<int, BooleanLattice> Lattice2;
+  Lattice2.insert({Key2, BooleanLattice(true)});
+
+  ASSERT_EQ(Lattice1.join(Lattice2), LatticeJoinEffect::Changed);
+  EXPECT_THAT(Lattice1, UnorderedElementsAre(Pair(Key1, BooleanLattice(true)),
+                                             Pair(Key2, BooleanLattice(true))));
+}
+
+TEST(MapLatticeTest, FindWorks) {
+  MapLattice<int, BooleanLattice> Lattice;
+  Lattice.insert({Key1, BooleanLattice(true)});
+  Lattice.insert({Key2, BooleanLattice(false)});
+
+  auto It = Lattice.find(Key1);
+  ASSERT_NE(It, Lattice.end());
+  EXPECT_EQ(It->second, BooleanLattice(true));
+
+  It = Lattice.find(Key2);
+  ASSERT_NE(It, Lattice.end());
+  EXPECT_EQ(It->second, BooleanLattice(false));
+}
+
+TEST(MapLatticeTest, ContainsWorks) {
+  MapLattice<int, BooleanLattice> Lattice;
+  Lattice.insert({Key1, BooleanLattice(true)});
+  EXPECT_TRUE(Lattice.contains(Key1));
+  EXPECT_FALSE(Lattice.contains(Key2));
+}
+} // namespace