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"
 
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 {}) {}
 
  // 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:
  using ValueMap = mlir::DenseMap<mlir::Value, ConstrainRefLatticeValue>;
  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(mlir::Value v, const ConstrainRefLatticeValue &rhs) {
    return valMap[v].setValue(rhs);
  }
 
  mlir::ChangeResult setValue(mlir::Value v, const ConstrainRef &ref) {
    return valMap[v].setValue(ConstrainRefLatticeValue(ref));
  }
 
  ConstrainRefLatticeValue getOrDefault(mlir::Value v) const;
 
  ConstrainRefLatticeValue getReturnValue(unsigned i) const;
 
  size_t size() const { return valMap.size(); }
 
  const ValueMap &getMap() const { return valMap; }
 
  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;
};
 
} // namespace llzk