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/SourceRef.h"
#include "llzk/Analysis/SparseAnalysis.h"
#include "llzk/Util/ErrorHelper.h"
 
#include <mlir/Analysis/DataFlow/DenseAnalysis.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(std::move(s)) {}
  explicit SourceRefLatticeValue(SourceRef r) : Base(ScalarTy {std::move(r)}) {}
  SourceRefLatticeValue() : Base(ScalarTy {}) {}
  SourceRefLatticeValue(const SourceRefLatticeValue &) = default;
  SourceRefLatticeValue(SourceRefLatticeValue &&) = default;
  SourceRefLatticeValue &operator=(const SourceRefLatticeValue &) = default;
  SourceRefLatticeValue &operator=(SourceRefLatticeValue &&) = default;
  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::FailureOr<llvm::DynamicAPInt> getConstantValue() const {
    if (isSingleValue()) {
      return getSingleValue().getConstantValue();
    }
    return mlir::failure();
  }

  mlir::FailureOr<llvm::DynamicAPInt> getConstantValue(const Field &field) const {
    if (isSingleValue()) {
      return getSingleValue().getConstantValue(field);
    }
    return mlir::failure();
  }

  using Base::getArrayShape;

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

  mlir::ChangeResult remove(const SourceRef &ref);

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

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

  mlir::ChangeResult write(
      const std::vector<SourceRefIndex> &indices, const SourceRefLatticeValue &rhs,
      bool joinWithExisting = false
  );

  mlir::FailureOr<std::pair<SourceRefLatticeValue, mlir::ChangeResult>>
  referenceMember(SymbolLookupResult<component::MemberDefOp> memberRef) const;

  mlir::FailureOr<std::pair<SourceRefLatticeValue, mlir::ChangeResult>>
  referencePodRecord(mlir::StringAttr recordName) const;

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

  mlir::ChangeResult replacePrefixesScalar(const TranslationMap &translation);

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

class SourceRefLattice : public dataflow::AbstractSparseLattice {
public:
  using LatticeValue = SourceRefLatticeValue;
  // 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 Ref2Val = mlir::DenseMap<SourceRef, mlir::DenseSet<ValueTy>>;
 
  /* Static utilities */

  static mlir::FailureOr<SourceRef> getSourceRef(mlir::Value val);
  static SourceRefLatticeValue getDefaultValue(ValueTy v);
 
  using AbstractSparseLattice::AbstractSparseLattice;
 
  mlir::ChangeResult join(const AbstractSparseLattice &rhs) override;
  mlir::ChangeResult meet(const AbstractSparseLattice &rhs) override;
  void print(mlir::raw_ostream &os) const override;
 
  const LatticeValue &getValue() const { return value; }
 
  mlir::ChangeResult setValue(const LatticeValue &newValue);
  mlir::ChangeResult setValue(const SourceRef &ref);
 
private:
  LatticeValue value;
};
 
} // namespace llzk
 
namespace llvm {
class raw_ostream;
 
raw_ostream &operator<<(raw_ostream &os, llvm::PointerUnion<mlir::Value, mlir::Operation *> ptr);
} // namespace llvm