AbstractLatticeValue.h

Go to the documentation of this file.

//===-- AbstractLatticeValue.h ----------------------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2025 Veridise Inc.
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include "llzk/Util/Debug.h"
#include "llzk/Util/ErrorHelper.h"
 
#include <mlir/Analysis/DataFlow/DeadCodeAnalysis.h>
#include <mlir/Support/LLVM.h>
 
#include <llvm/Support/Debug.h>
 
#include <concepts>
#include <type_traits>
#include <variant>
 
#define DEBUG_TYPE "llzk-abstract-lattice-value"
 
namespace llzk::dataflow {
 
template <typename Val>
concept ScalarLatticeValue = requires(Val lhs, Val rhs, mlir::raw_ostream &os) {
  // Require a form of print function
  { os << lhs } -> std::same_as<mlir::raw_ostream &>;
  // Require comparability
  { lhs == rhs } -> std::same_as<bool>;
  // Require the ability to combine two scalar values
  { lhs.join(rhs) } -> std::same_as<Val &>;
  // Require default constructable
  requires std::default_initializable<Val>;
};
 
template <typename Derived, ScalarLatticeValue ScalarTy> class AbstractLatticeValue {
  using ArrayTy = std::vector<std::unique_ptr<Derived>>;

  static ArrayTy constructArrayTy(const mlir::ArrayRef<int64_t> &shape) {
    size_t totalElem = 1;
    for (auto dim : shape) {
      totalElem *= dim;
    }
    ArrayTy arr(totalElem);
    for (auto it = arr.begin(); it != arr.end(); it++) {
      *it = std::make_unique<Derived>();
    }
    return arr;
  }
 
public:
  explicit AbstractLatticeValue(ScalarTy s) : value(s), arrayShape(std::nullopt) {}
  AbstractLatticeValue() : AbstractLatticeValue(ScalarTy()) {}
  explicit AbstractLatticeValue(const mlir::ArrayRef<int64_t> shape)
      : value(constructArrayTy(shape)), arrayShape(shape) {}
 
  AbstractLatticeValue(const AbstractLatticeValue &rhs) { *this = rhs; }
 
  // Enable copying by duplicating unique_ptrs and copying the contained values.
  AbstractLatticeValue &operator=(const AbstractLatticeValue &rhs) {
    copyArrayShape(rhs);
    if (rhs.isScalar()) {
      getValue() = rhs.getScalarValue();
    } else {
      // create an empty array of the same size
      getValue() = constructArrayTy(rhs.getArrayShape());
      auto &lhsArr = getArrayValue();
      auto &rhsArr = rhs.getArrayValue();
      for (unsigned i = 0; i < lhsArr.size(); i++) {
        // Recursive copy assignment of lattice values
        *lhsArr[i] = *rhsArr[i];
      }
    }
    return *this;
  }
 
  bool isScalar() const { return std::holds_alternative<ScalarTy>(value); }
  bool isSingleValue() const { return isScalar() && getScalarValue().size() == 1; }
  bool isArray() const { return std::holds_alternative<ArrayTy>(value); }
 
  const ScalarTy &getScalarValue() const {
    ensure(isScalar(), "not a scalar value");
    return std::get<ScalarTy>(value);
  }
 
  ScalarTy &getScalarValue() {
    ensure(isScalar(), "not a scalar value");
    return std::get<ScalarTy>(value);
  }
 
  const ArrayTy &getArrayValue() const {
    ensure(isArray(), "not an array value");
    return std::get<ArrayTy>(value);
  }
 
  ArrayTy &getArrayValue() {
    ensure(isArray(), "not an array value");
    return std::get<ArrayTy>(value);
  }

  const Derived &getElemFlatIdx(unsigned i) const {
    ensure(isArray(), "not an array value");
    auto &arr = getArrayValue();
    ensure(i < arr.size(), "index out of range");
    return *arr.at(i);
  }
 
  Derived &getElemFlatIdx(unsigned i) {
    ensure(isArray(), "not an array value");
    auto &arr = getArrayValue();
    ensure(i < arr.size(), "index out of range");
    return *arr.at(i);
  }
 
  size_t getArraySize() const { return getArrayValue().size(); }
 
  size_t getNumArrayDims() const { return getArrayShape().size(); }
 
  void print(mlir::raw_ostream &os) const {
    if (isScalar()) {
      os << getScalarValue();
    } else {
      os << "[ ";
      const auto &arr = getArrayValue();
      for (auto it = arr.begin(); it != arr.end();) {
        (*it)->print(os);
        it++;
        if (it != arr.end()) {
          os << ", ";
        } else {
          os << ' ';
        }
      }
      os << ']';
    }
  }

  ScalarTy foldToScalar() const {
    if (isScalar()) {
      return getScalarValue();
    }
 
    ScalarTy res;
    for (auto &val : getArrayValue()) {
      auto rhs = val->foldToScalar();
      res.join(rhs);
    }
    return res;
  }

  mlir::ChangeResult setValue(const AbstractLatticeValue &rhs) {
    if (*this == rhs) {
      return mlir::ChangeResult::NoChange;
    }
    *this = rhs;
    return mlir::ChangeResult::Change;
  }

  mlir::ChangeResult update(const Derived &rhs) {
    if (isScalar() && rhs.isScalar()) {
      return updateScalar(rhs.getScalarValue());
    } else if (isArray() && rhs.isArray() && getArraySize() == rhs.getArraySize()) {
      return updateArray(rhs.getArrayValue());
    } else {
      return foldAndUpdate(rhs);
    }
  }
 
  bool operator==(const AbstractLatticeValue &rhs) const {
    if (isScalar() && rhs.isScalar()) {
      return getScalarValue() == rhs.getScalarValue();
    } else if (isArray() && rhs.isArray() && getArraySize() == rhs.getArraySize()) {
      for (size_t i = 0; i < getArraySize(); i++) {
        if (getElemFlatIdx(i) != rhs.getElemFlatIdx(i)) {
          return false;
        }
      }
      return true;
    }
    return false;
  }
 
protected:
  std::variant<ScalarTy, ArrayTy> &getValue() { return value; }
 
  const std::vector<int64_t> &getArrayShape() const {
    ensure(arrayShape != std::nullopt, "not an array value");
    return arrayShape.value();
  }
 
  int64_t getArrayDim(unsigned i) const {
    const auto &arrShape = getArrayShape();
    ensure(i < arrShape.size(), "dimension index out of bounds");
    return arrShape.at(i);
  }
 
  void copyArrayShape(const AbstractLatticeValue &rhs) { arrayShape = rhs.arrayShape; }

  mlir::ChangeResult updateScalar(const ScalarTy &rhs) {
    auto lhs = getScalarValue();
    lhs.join(rhs);
    if (getScalarValue() == lhs) {
      return mlir::ChangeResult::NoChange;
    }
    getScalarValue() = lhs;
    return mlir::ChangeResult::Change;
  }

  mlir::ChangeResult updateArray(const ArrayTy &rhs) {
    mlir::ChangeResult res = mlir::ChangeResult::NoChange;
    auto &lhs = getArrayValue();
    for (size_t i = 0; i < getArraySize(); i++) {
      res |= lhs[i]->update(*rhs.at(i));
    }
    return res;
  }

  mlir::ChangeResult foldAndUpdate(const Derived &rhs) {
    auto folded = foldToScalar();
    auto rhsScalar = rhs.foldToScalar();
    folded.join(rhsScalar);
    if (isScalar() && getScalarValue() == folded) {
      return mlir::ChangeResult::NoChange;
    }
    getValue() = folded;
    return mlir::ChangeResult::Change;
  }
 
private:
  std::variant<ScalarTy, ArrayTy> value;
  std::optional<std::vector<int64_t>> arrayShape;
};
 
template <typename Derived, ScalarLatticeValue ScalarTy>
mlir::raw_ostream &
operator<<(mlir::raw_ostream &os, const AbstractLatticeValue<Derived, ScalarTy> &v) {
  v.print(os);
  return os;
}
 
} // namespace llzk::dataflow
 
#undef DEBUG_TYPE