ConstrainRefLattice.h

Go to the documentation of this file.

//===-- ConstrainRefLattice.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/Analysis/AbstractLatticeValue.h"
#include "llzk/Analysis/ConstrainRef.h"
#include "llzk/Analysis/DenseAnalysis.h"
#include "llzk/Util/ErrorHelper.h"
 
#include <llvm/ADT/PointerUnion.h>
 
namespace llzk {
 
class ConstrainRefLatticeValue;
using TranslationMap =
    std::unordered_map<ConstrainRef, ConstrainRefLatticeValue, ConstrainRef::Hash>;

class ConstrainRefLatticeValue
    : public dataflow::AbstractLatticeValue<ConstrainRefLatticeValue, ConstrainRefSet> {
  using Base = dataflow::AbstractLatticeValue<ConstrainRefLatticeValue, ConstrainRefSet>;
  using ScalarTy = ConstrainRefSet;
  using ArrayTy = std::vector<std::unique_ptr<ConstrainRefLatticeValue>>;
 
public:
  explicit ConstrainRefLatticeValue(ScalarTy s) : Base(s) {}
  explicit ConstrainRefLatticeValue(ConstrainRef r) : Base(ScalarTy {r}) {}
  ConstrainRefLatticeValue() : Base(ScalarTy {}) {}
  virtual ~ConstrainRefLatticeValue() = default;
 
  // Create an empty array of the given shape.
  explicit ConstrainRefLatticeValue(mlir::ArrayRef<int64_t> shape) : Base(shape) {}
 
  const ConstrainRef &getSingleValue() const {
    ensure(isSingleValue(), "not a single value");
    return *getScalarValue().begin();
  }

  mlir::ChangeResult insert(const ConstrainRef &rhs);

  std::pair<ConstrainRefLatticeValue, mlir::ChangeResult>
  translate(const TranslationMap &translation) const;

  std::pair<ConstrainRefLatticeValue, mlir::ChangeResult>
  referenceField(SymbolLookupResult<component::FieldDefOp> fieldRef) const;

  std::pair<ConstrainRefLatticeValue, mlir::ChangeResult>
  extract(const std::vector<ConstrainRefIndex> &indices) const;
 
protected:
  mlir::ChangeResult translateScalar(const TranslationMap &translation);

  virtual std::pair<ConstrainRefLatticeValue, mlir::ChangeResult>
  elementwiseTransform(llvm::function_ref<ConstrainRef(const ConstrainRef &)> transform) const;
};

class ConstrainRefLattice : public dataflow::AbstractDenseLattice {
public:
  // mlir::Value is used for read-like operations that create references in their results,
  // mlir::Operation* is used for write-like operations that reference values as their destinations
  using ValueTy = llvm::PointerUnion<mlir::Value, mlir::Operation *>;
  using ValueMap = mlir::DenseMap<ValueTy, ConstrainRefLatticeValue>;
  // Used to lookup MLIR values/operations from a given ConstrainRef (all values that a ref is
  // referenced by)
  using ValueSet = mlir::DenseSet<ValueTy>;
  using Ref2Val = mlir::DenseMap<ConstrainRef, mlir::DenseSet<ValueTy>>;
  using AbstractDenseLattice::AbstractDenseLattice;
 
  /* Static utilities */

  static mlir::FailureOr<ConstrainRef> getSourceRef(mlir::Value val);
 
  /* Required methods */

  mlir::ChangeResult join(const AbstractDenseLattice &rhs) override {
    if (auto *r = dynamic_cast<const ConstrainRefLattice *>(&rhs)) {
      return setValues(r->valMap);
    }
    llvm::report_fatal_error("invalid join lattice type");
    return mlir::ChangeResult::NoChange;
  }

  virtual mlir::ChangeResult meet(const AbstractDenseLattice &rhs) override {
    llvm::report_fatal_error("meet operation is not supported for ConstrainRefLattice");
    return mlir::ChangeResult::NoChange;
  }
 
  void print(mlir::raw_ostream &os) const override;
 
  /* Update utility methods */
 
  mlir::ChangeResult setValues(const ValueMap &rhs);
 
  mlir::ChangeResult setValue(ValueTy v, const ConstrainRefLatticeValue &rhs);
 
  mlir::ChangeResult setValue(ValueTy v, const ConstrainRef &ref);
 
  ConstrainRefLatticeValue getOrDefault(ValueTy v) const;
 
  ConstrainRefLatticeValue getReturnValue(unsigned i) const;
 
  ValueSet lookupValues(const ConstrainRef &r) const;
 
  size_t size() const { return valMap.size(); }
 
  const ValueMap &getMap() const { return valMap; }
 
  const Ref2Val &getRef2Val() const { return refMap; }
 
  ValueMap::iterator begin() { return valMap.begin(); }
  ValueMap::iterator end() { return valMap.end(); }
  ValueMap::const_iterator begin() const { return valMap.begin(); }
  ValueMap::const_iterator end() const { return valMap.end(); }
 
  friend mlir::raw_ostream &operator<<(mlir::raw_ostream &os, const ConstrainRefLattice &v);
 
private:
  ValueMap valMap;
  Ref2Val refMap;
};
 
} // namespace llzk
 
namespace llvm {
class raw_ostream;
 
raw_ostream &operator<<(raw_ostream &os, llvm::PointerUnion<mlir::Value, mlir::Operation *> ptr);
} // namespace llvm