From 5ee9ef85190cff09c4d3718d8683af1900a33348 Mon Sep 17 00:00:00 2001
From: Sander de Smalen <sander.desmalen@arm.com>
Date: Fri, 6 Nov 2020 07:56:55 +0000
Subject: [PATCH] [TypeSize] Extend UnivariateLinearPolyBase with
 getWithIncrement/Decrement methods

This patch adds getWithIncrement/getWithDecrement methods to
ElementCount and TypeSize to allow:

  TypeSize::getFixed(8).getWithIncrement(8)     <=> TypeSize::getFixed(16)
  TypeSize::getFixed(16).getWithDecrement(8)    <=> TypeSize::getFixed(8)
  TypeSize::getScalable(8).getWithIncrement(8)  <=> TypeSize::getScalable(16)
  TypeSize::getScalable(16).getWithDecrement(8) <=> TypeSize::getScalable(8)

This patch implements parts of the POC in D90342.

Reviewed By: ctetreau, dmgreen

Differential Revision: https://reviews.llvm.org/D90713
---
 llvm/include/llvm/Support/TypeSize.h          | 23 ++++++++++++++++++++++-
 llvm/unittests/Support/LinearPolyBaseTest.cpp | 10 ++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/Support/TypeSize.h b/llvm/include/llvm/Support/TypeSize.h
index be40be2..d277aff 100644
--- a/llvm/include/llvm/Support/TypeSize.h
+++ b/llvm/include/llvm/Support/TypeSize.h
@@ -169,7 +169,7 @@ protected:
   ScalarTy Value;         // The value at the univeriate dimension.
   unsigned UnivariateDim; // The univeriate dimension.
 
-  UnivariateLinearPolyBase(ScalarTy &Val, unsigned UnivariateDim)
+  UnivariateLinearPolyBase(ScalarTy Val, unsigned UnivariateDim)
       : Value(Val), UnivariateDim(UnivariateDim) {
     assert(UnivariateDim < Dimensions && "Dimension out of range");
   }
@@ -229,6 +229,18 @@ public:
   ScalarTy getValue(unsigned Dim) const {
     return Dim == UnivariateDim ? Value : 0;
   }
+
+  /// Add \p RHS to the value at the univariate dimension.
+  LeafTy getWithIncrement(ScalarTy RHS) {
+    return static_cast<LeafTy>(
+        UnivariateLinearPolyBase(Value + RHS, UnivariateDim));
+  }
+
+  /// Subtract \p RHS from the value at the univariate dimension.
+  LeafTy getWithDecrement(ScalarTy RHS) {
+    return static_cast<LeafTy>(
+        UnivariateLinearPolyBase(Value - RHS, UnivariateDim));
+  }
 };
 
 
@@ -247,6 +259,11 @@ public:
 /// fixed-sized or it is scalable-sized, but it cannot be both.
 template <typename LeafTy>
 class LinearPolySize : public UnivariateLinearPolyBase<LeafTy> {
+  // Make the parent class a friend, so that it can access the protected
+  // conversion/copy-constructor for UnivariatePolyBase<LeafTy> ->
+  // LinearPolySize<LeafTy>.
+  friend class UnivariateLinearPolyBase<LeafTy>;
+
 public:
   using ScalarTy = typename UnivariateLinearPolyBase<LeafTy>::ScalarTy;
   enum Dims : unsigned { FixedDim = 0, ScalableDim = 1 };
@@ -255,7 +272,11 @@ protected:
   LinearPolySize(ScalarTy MinVal, Dims D)
       : UnivariateLinearPolyBase<LeafTy>(MinVal, D) {}
 
+  LinearPolySize(const UnivariateLinearPolyBase<LeafTy> &V)
+      : UnivariateLinearPolyBase<LeafTy>(V) {}
+
 public:
+
   static LeafTy getFixed(ScalarTy MinVal) {
     return static_cast<LeafTy>(LinearPolySize(MinVal, FixedDim));
   }
diff --git a/llvm/unittests/Support/LinearPolyBaseTest.cpp b/llvm/unittests/Support/LinearPolyBaseTest.cpp
index 407e9b6..6942ea2 100644
--- a/llvm/unittests/Support/LinearPolyBaseTest.cpp
+++ b/llvm/unittests/Support/LinearPolyBaseTest.cpp
@@ -140,6 +140,11 @@ TEST(UnivariateLinearPolyBase, Univariate3D_Add) {
   Univariate3D X(42, 0);
   X += Univariate3D(42, 0);
   EXPECT_EQ(X, Univariate3D(84, 0));
+
+  // Test 'getWithIncrement' method
+  EXPECT_EQ(Univariate3D(42, 0).getWithIncrement(1), Univariate3D(43, 0));
+  EXPECT_EQ(Univariate3D(42, 1).getWithIncrement(2), Univariate3D(44, 1));
+  EXPECT_EQ(Univariate3D(42, 2).getWithIncrement(3), Univariate3D(45, 2));
 }
 
 TEST(UnivariateLinearPolyBase, Univariate3D_Sub) {
@@ -153,6 +158,11 @@ TEST(UnivariateLinearPolyBase, Univariate3D_Sub) {
   Univariate3D X(84, 0);
   X -= Univariate3D(42, 0);
   EXPECT_EQ(X, Univariate3D(42, 0));
+
+  // Test 'getWithDecrement' method
+  EXPECT_EQ(Univariate3D(43, 0).getWithDecrement(1), Univariate3D(42, 0));
+  EXPECT_EQ(Univariate3D(44, 1).getWithDecrement(2), Univariate3D(42, 1));
+  EXPECT_EQ(Univariate3D(45, 2).getWithDecrement(3), Univariate3D(42, 2));
 }
 
 TEST(UnivariateLinearPolyBase, Univariate3D_Scale) {
-- 
2.7.4