reflect.hpp

Go to the documentation of this file.

#pragma once
 
#include <algorithm>
#include <boost/describe.hpp>
#include <cstdint>
#include <map>
#include <ranges>
#include <span>
#include <spirv/unified1/spirv.hpp>
#include <string>
#include <vector>
 
#include "wren/utils/binary_reader.hpp"
 
namespace r = std::ranges;
namespace rv = std::views;
 
namespace spv {
BOOST_DESCRIBE_ENUM(StorageClass, StorageClassAtomicCounter)
BOOST_DESCRIBE_ENUM(Op, OpCapability, OpDecorate, OpEntryPoint, OpName,
                    OpTypeInt, OpTypeFloat, OpTypeVector, OpTypePointer,
                    OpVariable)
}  // namespace spv
 
namespace wren::reflect {
 
struct EntryPoint {
  std::string name;
  spv::ExecutionModel execution_model;
  uint32_t entry_point;
 
  std::vector<uint32_t> forward_references;
};
 
struct Variable {
  std::string name;
  spv::StorageClass storage_class;
  uint32_t location;
  uint32_t width = UINT32_MAX;
  uint32_t format = UINT32_MAX;  // TODO(tyler) This needs updating for textures
 
  auto operator==(const Variable&) const -> bool = default;
  auto operator!=(const Variable&) const -> bool = default;
};
BOOST_DESCRIBE_STRUCT(Variable, (),
                      (name, storage_class, location, width, format))
 
BOOST_DEFINE_ENUM_CLASS(TypeType, Float, Vector);
 
struct SimpleType {
  uint32_t width;
};
 
struct Vector {
  SimpleType type;
  uint32_t count;
};
 
struct Annotation {
  spv::Decoration decoration = spv::Decoration::DecorationMax;
  std::optional<uint32_t> location;
  std::optional<uint32_t> binding;
  std::optional<uint32_t> descriptor_set;
};
 
struct Binding {
  uint32_t binding;
  uint32_t descriptor_type;
  uint32_t count;
  auto operator==(const Binding&) const -> bool = default;
};
 
BOOST_DESCRIBE_STRUCT(Binding, (), (binding, descriptor_type, count));
 
struct DescriptorSet {
  std::vector<Binding> bindings;
  auto operator==(const DescriptorSet&) const -> bool = default;
};
BOOST_DESCRIBE_STRUCT(DescriptorSet, (), (bindings));
 
class Reflect {
 public:
  Reflect() = default;
  Reflect(const std::span<const std::byte>& spirv);
 
  [[nodiscard]] auto magic_number() const { return magic_number_; }
 
  [[nodiscard]] auto entry_points() const { return entry_points_; }
 
  [[nodiscard]] auto input_variables() const -> std::vector<Variable> {
    return variables_ | rv::filter([](const auto& id_v) {
             const auto& [_, v] = id_v;
             return v.storage_class == spv::StorageClass::StorageClassInput;
           }) |
           rv::transform([](const auto& id_v) { return id_v.second; }) |
           r::to<std::vector>();
  }
 
  [[nodiscard]] auto input_variables_for_entry_point(
      spv::ExecutionModel entry_point_type) const -> std::vector<Variable> {
    const auto& entry_point =
        r::find_if(entry_points_, [&](const EntryPoint& e) {
          return e.execution_model == entry_point_type;
        });
 
    return variables_ | rv::filter([entry_point](const auto& id_v) {
             const auto& [id, v] = id_v;
             return r::contains(entry_point->forward_references, id) &&
                    v.storage_class == spv::StorageClass::StorageClassInput;
           }) |
           rv::transform([](const auto& id_v) { return id_v.second; }) |
           r::to<std::vector>();
  }
 
  [[nodiscard]] auto descriptor_sets() const -> std::vector<DescriptorSet> {
    return descriptor_sets_ |
           rv::transform([](const auto& ds) { return ds.second; }) |
           r::to<std::vector>();
  }
 
 private:
  void resolve();
  auto parse_op_code(utils::BinaryReader& reader, uint32_t wordcount,
                     const spv::Op& opcode) -> uint32_t;
 
  uint32_t magic_number_{};
  uint32_t version_{};
  uint32_t generator_{};
  uint32_t bound_{};
 
  std::vector<EntryPoint> entry_points_;
  std::map<uint32_t, Variable> variables_;
  std::map<uint32_t, DescriptorSet> descriptor_sets_;
 
  std::map<uint32_t, std::vector<Annotation>> annotations_;
  std::map<uint32_t, SimpleType> types_;
 
  std::map<uint32_t, std::string> names_;
};
 
}  // namespace wren::reflect