SourceRefLattice.h

Go to the documentation of this file.

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

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

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

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

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

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

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

class SourceRefLattice : 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, SourceRefLatticeValue>;
  // Used to lookup MLIR values/operations from a given SourceRef (all values that a ref is
  // referenced by)
  using ValueSet = mlir::DenseSet<ValueTy>;
  using Ref2Val = mlir::DenseMap<SourceRef, mlir::DenseSet<ValueTy>>;
  using AbstractDenseLattice::AbstractDenseLattice;
 
  /* Static utilities */

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

  mlir::ChangeResult join(const AbstractDenseLattice &rhs) override {
    return setValues(static_cast<const SourceRefLattice &>(rhs).valMap);
  }

  mlir::ChangeResult meet(const AbstractDenseLattice & /*rhs*/) override {
    llvm::report_fatal_error("meet operation is not supported for SourceRefLattice");
    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 SourceRefLatticeValue &rhs);
 
  mlir::ChangeResult setValue(ValueTy v, const SourceRef &ref);
 
  SourceRefLatticeValue getOrDefault(ValueTy v) const;
 
  SourceRefLatticeValue getReturnValue(unsigned i) const;
 
  ValueSet lookupValues(const SourceRef &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 SourceRefLattice &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