[analysis] Implement an array lattice (#6057)

The elements of `Array<L, N>` lattice are arrays of length `N` of elements of
`L`, compared pairwise with each other. This lattice is a concrete
implementation of what would be written L^N with pen and paper.

3 files changed, 280 insertions, 9 deletions

diff --git a/src/analysis/lattices/array.h b/src/analysis/lattices/array.h
new file mode 100644
index 000000000..3426a3a77
--- /dev/null
+++ b/src/analysis/lattices/array.h
@@ -0,0 +1,122 @@
+/*
+ * Copyright 2023 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef wasm_analysis_lattices_array_h
+#define wasm_analysis_lattices_array_h
+
+#include <array>
+#include <utility>
+
+#include "../lattice.h"
+#include "bool.h"
+#include "flat.h"
+
+namespace wasm::analysis {
+
+// A lattice whose elements are N-tuples of elements of L. Also written as L^N.
+template<Lattice L, size_t N> struct Array {
+  using Element = std::array<typename L::Element, N>;
+
+  L lattice;
+
+  Array(L&& lattice) : lattice(std::move(lattice)) {}
+
+private:
+  // Use a template parameter pack to generate N copies of
+  // `lattice.getBottom()`. TODO: Use C++20 lambda template parameters instead
+  // of a separate helper function.
+  template<size_t... I>
+  Element getBottomImpl(std::index_sequence<I...>) const noexcept {
+    return {((void)I, lattice.getBottom())...};
+  }
+  template<size_t... I>
+  Element getTopImpl(std::index_sequence<I...>) const noexcept {
+    return {((void)I, lattice.getTop())...};
+  }
+
+public:
+  Element getBottom() const noexcept {
+    return getBottomImpl(std::make_index_sequence<N>());
+  }
+
+  Element getTop() const noexcept
+#if __cplusplus >= 202002L
+    requires FullLattice<L>
+#endif
+  {
+    return getTopImpl(std::make_index_sequence<N>());
+  }
+
+  // `a` <= `b` if all their elements are pairwise <=, etc. Unless we determine
+  // that there is no relation, we must check all the elements.
+  LatticeComparison compare(const Element& a, const Element& b) const noexcept {
+    auto result = EQUAL;
+    for (size_t i = 0; i < N; ++i) {
+      switch (lattice.compare(a[i], b[i])) {
+        case NO_RELATION:
+          return NO_RELATION;
+        case EQUAL:
+          continue;
+        case LESS:
+          if (result == GREATER) {
+            // Cannot be both less and greater.
+            return NO_RELATION;
+          }
+          result = LESS;
+          continue;
+        case GREATER:
+          if (result == LESS) {
+            // Cannot be both greater and less.
+            return NO_RELATION;
+          }
+          result = GREATER;
+          continue;
+      }
+    }
+    return result;
+  }
+
+  // Pairwise join on the elements.
+  bool join(Element& joinee, const Element& joiner) const noexcept {
+    bool result = false;
+    for (size_t i = 0; i < N; ++i) {
+      result |= lattice.join(joinee[i], joiner[i]);
+    }
+    return result;
+  }
+
+  // Pairwise meet on the elements.
+  bool meet(Element& meetee, const Element& meeter) const noexcept
+#if __cplusplus >= 202002L
+    requires FullLattice<L>
+#endif
+  {
+    bool result = false;
+    for (size_t i = 0; i < N; ++i) {
+      result |= lattice.meet(meetee[i], meeter[i]);
+    }
+    return result;
+  }
+};
+
+#if __cplusplus >= 202002L
+static_assert(FullLattice<Array<Bool, 1>>);
+static_assert(Lattice<Array<Flat<bool>, 1>>);
+#endif
+
+} // namespace wasm::analysis
+
+#endif // wasm_analysis_lattices_array_h
diff --git a/src/tools/wasm-fuzz-lattices.cpp b/src/tools/wasm-fuzz-lattices.cpp
index abc855467..f073aa724 100644
--- a/src/tools/wasm-fuzz-lattices.cpp
+++ b/src/tools/wasm-fuzz-lattices.cpp
@@ -21,6 +21,7 @@
 #include <variant>
 
 #include "analysis/lattice.h"
+#include "analysis/lattices/array.h"
 #include "analysis/lattices/bool.h"
 #include "analysis/lattices/flat.h"
 #include "analysis/lattices/int.h"
@@ -147,28 +148,36 @@ static_assert(FullLattice<RandomFullLattice>);
 static_assert(Lattice<RandomLattice>);
 #endif
 
+using ArrayFullLattice = analysis::Array<RandomFullLattice, 2>;
+using ArrayLattice = analysis::Array<RandomLattice, 2>;
+
 struct RandomFullLattice::L
-  : std::variant<Bool, UInt32, Inverted<RandomFullLattice>> {};
+  : std::variant<Bool, UInt32, Inverted<RandomFullLattice>, ArrayFullLattice> {
+};
 
 struct RandomFullLattice::ElementImpl
   : std::variant<typename Bool::Element,
                  typename UInt32::Element,
-                 typename Inverted<RandomFullLattice>::Element> {};
+                 typename Inverted<RandomFullLattice>::Element,
+                 typename ArrayFullLattice::Element> {};
 
-struct RandomLattice::L
-  : std::variant<RandomFullLattice, Flat<uint32_t>, Lift<RandomLattice>> {};
+struct RandomLattice::L : std::variant<RandomFullLattice,
+                                       Flat<uint32_t>,
+                                       Lift<RandomLattice>,
+                                       ArrayLattice> {};
 
 struct RandomLattice::ElementImpl
   : std::variant<typename RandomFullLattice::Element,
                  typename Flat<uint32_t>::Element,
-                 typename Lift<RandomLattice>::Element> {};
+                 typename Lift<RandomLattice>::Element,
+                 typename ArrayLattice::Element> {};
 
 RandomFullLattice::RandomFullLattice(Random& rand,
                                      size_t depth,
                                      std::optional<uint32_t> maybePick)
   : rand(rand) {
   // TODO: Limit the depth once we get lattices with more fan-out.
-  uint32_t pick = maybePick ? *maybePick : rand.upTo(3);
+  uint32_t pick = maybePick ? *maybePick : rand.upTo(4);
   switch (pick) {
     case 0:
       lattice = std::make_unique<L>(L{Bool{}});
@@ -180,25 +189,34 @@ RandomFullLattice::RandomFullLattice(Random& rand,
       lattice =
         std::make_unique<L>(L{Inverted{RandomFullLattice{rand, depth + 1}}});
       return;
+    case 3:
+      lattice = std::make_unique<L>(
+        L{ArrayFullLattice{RandomFullLattice{rand, depth + 1}}});
+      return;
   }
   WASM_UNREACHABLE("unexpected pick");
 }
 
 RandomLattice::RandomLattice(Random& rand, size_t depth) : rand(rand) {
   // TODO: Limit the depth once we get lattices with more fan-out.
-  uint32_t pick = rand.upTo(5);
+  uint32_t pick = rand.upTo(7);
   switch (pick) {
     case 0:
     case 1:
     case 2:
+    case 3:
       lattice = std::make_unique<L>(L{RandomFullLattice{rand, depth, pick}});
       return;
-    case 3:
+    case 4:
       lattice = std::make_unique<L>(L{Flat<uint32_t>{}});
       return;
-    case 4:
+    case 5:
       lattice = std::make_unique<L>(L{Lift{RandomLattice{rand, depth + 1}}});
       return;
+    case 6:
+      lattice =
+        std::make_unique<L>(L{ArrayLattice{RandomLattice{rand, depth + 1}}});
+      return;
   }
   WASM_UNREACHABLE("unexpected pick");
 }
@@ -213,6 +231,10 @@ RandomFullLattice::Element RandomFullLattice::makeElement() const noexcept {
   if (const auto* l = std::get_if<Inverted<RandomFullLattice>>(lattice.get())) {
     return ElementImpl{l->lattice.makeElement()};
   }
+  if (const auto* l = std::get_if<ArrayFullLattice>(lattice.get())) {
+    return ElementImpl{typename ArrayFullLattice::Element{
+      l->lattice.makeElement(), l->lattice.makeElement()}};
+  }
   WASM_UNREACHABLE("unexpected lattice");
 }
 
@@ -235,6 +257,10 @@ RandomLattice::Element RandomLattice::makeElement() const noexcept {
     return ElementImpl{rand.oneIn(4) ? l->getBottom()
                                      : l->get(l->lattice.makeElement())};
   }
+  if (const auto* l = std::get_if<ArrayLattice>(lattice.get())) {
+    return ElementImpl{typename ArrayLattice::Element{
+      l->lattice.makeElement(), l->lattice.makeElement()}};
+  }
   WASM_UNREACHABLE("unexpected lattice");
 }
 
@@ -260,6 +286,13 @@ void printFullElement(std::ostream& os,
     printFullElement(os, *e, depth + 1);
     indent(os, depth);
     os << ")\n";
+  } else if (const auto* e =
+               std::get_if<typename ArrayFullLattice::Element>(&*elem)) {
+    os << "Array[\n";
+    printFullElement(os, e->front(), depth + 1);
+    printFullElement(os, e->back(), depth + 1);
+    indent(os, depth);
+    os << "]\n";
   }
 }
 
@@ -292,6 +325,13 @@ void printElement(std::ostream& os,
       indent(os, depth);
       os << ")\n";
     }
+  } else if (const auto* e =
+               std::get_if<typename ArrayLattice::Element>(&*elem)) {
+    os << "Array[\n";
+    printElement(os, e->front(), depth + 1);
+    printElement(os, e->back(), depth + 1);
+    indent(os, depth);
+    os << ")\n";
   }
 }
 
diff --git a/test/gtest/lattices.cpp b/test/gtest/lattices.cpp
index e201c56b8..50adffc7e 100644
--- a/test/gtest/lattices.cpp
+++ b/test/gtest/lattices.cpp
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include "analysis/lattices/array.h"
 #include "analysis/lattices/bool.h"
 #include "analysis/lattices/flat.h"
 #include "analysis/lattices/int.h"
@@ -374,3 +375,111 @@ TEST(LiftLattice, Join) {
   EXPECT_FALSE(lift.join(elem, bot));
   EXPECT_EQ(elem, top);
 }
+
+TEST(ArrayLattice, GetBottom) {
+  analysis::Array<analysis::Bool, 2> array{analysis::Bool{}};
+  EXPECT_EQ(array.getBottom(), (std::array<bool, 2>{false, false}));
+}
+
+TEST(ArrayLattice, GetTop) {
+  analysis::Array<analysis::Bool, 2> array{analysis::Bool{}};
+  EXPECT_EQ(array.getTop(), (std::array<bool, 2>{true, true}));
+}
+
+TEST(ArrayLattice, Compare) {
+  analysis::Array<analysis::Bool, 2> array{analysis::Bool{}};
+  std::array<bool, 2> ff{false, false};
+  std::array<bool, 2> ft{false, true};
+  std::array<bool, 2> tf{true, false};
+  std::array<bool, 2> tt{true, true};
+
+  EXPECT_EQ(array.compare(ff, ff), analysis::EQUAL);
+  EXPECT_EQ(array.compare(ff, ft), analysis::LESS);
+  EXPECT_EQ(array.compare(ff, tf), analysis::LESS);
+  EXPECT_EQ(array.compare(ff, tt), analysis::LESS);
+
+  EXPECT_EQ(array.compare(ft, ff), analysis::GREATER);
+  EXPECT_EQ(array.compare(ft, ft), analysis::EQUAL);
+  EXPECT_EQ(array.compare(ft, tf), analysis::NO_RELATION);
+  EXPECT_EQ(array.compare(ft, tt), analysis::LESS);
+
+  EXPECT_EQ(array.compare(tf, ff), analysis::GREATER);
+  EXPECT_EQ(array.compare(tf, ft), analysis::NO_RELATION);
+  EXPECT_EQ(array.compare(tf, tf), analysis::EQUAL);
+  EXPECT_EQ(array.compare(tf, tt), analysis::LESS);
+
+  EXPECT_EQ(array.compare(tt, ff), analysis::GREATER);
+  EXPECT_EQ(array.compare(tt, ft), analysis::GREATER);
+  EXPECT_EQ(array.compare(tt, tf), analysis::GREATER);
+  EXPECT_EQ(array.compare(tt, tt), analysis::EQUAL);
+}
+
+TEST(ArrayLattice, Join) {
+  analysis::Array<analysis::Bool, 2> array{analysis::Bool{}};
+  auto ff = []() { return std::array<bool, 2>{false, false}; };
+  auto ft = []() { return std::array<bool, 2>{false, true}; };
+  auto tf = []() { return std::array<bool, 2>{true, false}; };
+  auto tt = []() { return std::array<bool, 2>{true, true}; };
+
+  auto test =
+    [&](auto& makeJoinee, auto& makeJoiner, bool modified, auto& makeExpected) {
+      auto joinee = makeJoinee();
+      EXPECT_EQ(array.join(joinee, makeJoiner()), modified);
+      EXPECT_EQ(joinee, makeExpected());
+    };
+
+  test(ff, ff, false, ff);
+  test(ff, ft, true, ft);
+  test(ff, tf, true, tf);
+  test(ff, tt, true, tt);
+
+  test(ft, ff, false, ft);
+  test(ft, ft, false, ft);
+  test(ft, tf, true, tt);
+  test(ft, tt, true, tt);
+
+  test(tf, ff, false, tf);
+  test(tf, ft, true, tt);
+  test(tf, tf, false, tf);
+  test(tf, tt, true, tt);
+
+  test(tt, ff, false, tt);
+  test(tt, ft, false, tt);
+  test(tt, tf, false, tt);
+  test(tt, tt, false, tt);
+}
+
+TEST(ArrayLattice, Meet) {
+  analysis::Array<analysis::Bool, 2> array{analysis::Bool{}};
+  auto ff = []() { return std::array<bool, 2>{false, false}; };
+  auto ft = []() { return std::array<bool, 2>{false, true}; };
+  auto tf = []() { return std::array<bool, 2>{true, false}; };
+  auto tt = []() { return std::array<bool, 2>{true, true}; };
+
+  auto test =
+    [&](auto& makeMeetee, auto& makeMeeter, bool modified, auto& makeExpected) {
+      auto meetee = makeMeetee();
+      EXPECT_EQ(array.meet(meetee, makeMeeter()), modified);
+      EXPECT_EQ(meetee, makeExpected());
+    };
+
+  test(ff, ff, false, ff);
+  test(ff, ft, false, ff);
+  test(ff, tf, false, ff);
+  test(ff, tt, false, ff);
+
+  test(ft, ff, true, ff);
+  test(ft, ft, false, ft);
+  test(ft, tf, true, ff);
+  test(ft, tt, false, ft);
+
+  test(tf, ff, true, ff);
+  test(tf, ft, true, ff);
+  test(tf, tf, false, tf);
+  test(tf, tt, false, tf);
+
+  test(tt, ff, true, ff);
+  test(tt, ft, true, ft);
+  test(tt, tf, true, tf);
+  test(tt, tt, false, tt);
+}