voxels.hpp

#pragma once

#include <torch/extension.h>
#include "utils.hpp"

namespace metapedia {

using namespace torch::indexing;

struct EmptyVoxel {};

struct ColorVoxel {
  uint32_t rgba;

  ColorVoxel(uint8_t r, uint8_t g, uint8_t b, uint8_t a)
      : rgba(pack_bytes(r, g, b, a)) {}
  ColorVoxel(int8_t r, int8_t g, int8_t b) : ColorVoxel(r, g, b, 255) {}
};

struct ThreeSliceSpec {
  Slice zero;
  Slice one;
  Slice two;

  auto index_into(torch::Tensor &tensor) {
    return tensor.index({zero, one, two});
  }
};

using VoxelDef = std::variant<EmptyVoxel, ColorVoxel>;

struct VoxelConfig {
  using Key = int32_t;
  std::unordered_map<Key, VoxelDef> defs;
};

using VoxelTensor = torch::Tensor;

struct VoxelMesh {
  std::vector<float> positions;
  std::vector<float> normals;
  std::vector<uint8_t> colors;
  std::vector<uint32_t> triangles;

  auto vertex_count() const {
    return positions.size() / 3;
  }

  auto triangle_count() const {
    return triangles.size() / 3;
  }
};

auto compute_adjacency_mask(const VoxelTensor& tensor) {
  assert(tensor.dim() == 3);

  auto w = tensor.size(0);
  auto h = tensor.size(1);
  auto d = tensor.size(2);
  auto sl = [&](int x_0, int x_1, int y_0, int y_1, int z_0, int z_1) {
    x_0 = x_0 < 0 ? w + x_0 : x_0;
    y_0 = y_0 < 0 ? h + y_0 : y_0;
    z_0 = z_0 < 0 ? d + z_0 : z_0;
    x_1 = x_1 <= 0 ? w + x_1 : x_1;
    y_1 = y_1 <= 0 ? h + y_1 : y_1;
    z_1 = z_1 <= 0 ? d + z_1 : z_1;
    return ThreeSliceSpec(
      {Slice(x_0, x_1, 1), Slice(y_0, y_1, 1), Slice(z_0, z_1, 1)}
    );
  };

  auto update = [](torch::Tensor& tensor, ThreeSliceSpec slice, torch::Tensor input) {
    auto merged = torch::bitwise_or(tensor.index({slice.zero, slice.one, slice.two}), input);
    tensor.index_put_({slice.zero, slice.one, slice.two}, merged);
  };

  auto empty = (tensor == 0) * ((int32_t)0xFFFFFFFF);
  auto adjacencies = torch::zeros_like(tensor);

  constexpr auto x_neg_mask = static_cast<int>(DirMask::X_NEG);
  constexpr auto x_pos_mask = static_cast<int>(DirMask::X_POS);
  constexpr auto y_neg_mask = static_cast<int>(DirMask::Y_NEG);
  constexpr auto y_pos_mask = static_cast<int>(DirMask::Y_POS);
  constexpr auto z_neg_mask = static_cast<int>(DirMask::Z_NEG);
  constexpr auto z_pos_mask = static_cast<int>(DirMask::Z_POS);

  update(
    adjacencies,
    sl(0, 1, 0, h, 0, d),
    torch::bitwise_and(torch::bitwise_not(
      sl(0, 1, 0, h, 0, d).index_into(empty)), 
      x_neg_mask));

  update(
    adjacencies,
    sl(-1, w, 0, h, 0, d),
    torch::bitwise_and(torch::bitwise_not(
      sl(-1, w, 0, h, 0, d).index_into(empty)),
      x_pos_mask));


  update(
    adjacencies,
    sl(0, w, 0, 1, 0, d),
    torch::bitwise_and(torch::bitwise_not(
      sl(0, w, 0, 1, 0, d).index_into(empty)),
      y_neg_mask
    ));

  update(
      adjacencies,
      sl(0, w, -1, h, 0, d),
      torch::bitwise_and(torch::bitwise_not(
        sl(0, w, -1, h, 0, d).index_into(empty)),
        y_pos_mask
      ));

  update(
      adjacencies,
      sl(0, w, 0, h, 0, 1),
      torch::bitwise_and(torch::bitwise_not(
        sl(0, w, 0, h, 0, 1).index_into(empty)),
        z_neg_mask
      ));

  update(
      adjacencies,
      sl(0, w, 0, h, -1, d),
      torch::bitwise_and(torch::bitwise_not(
        sl(0, w, 0, h, -1, d).index_into(empty)),
        z_pos_mask
      ));

  // Initialize the adjacency tensor interior values.
  auto shift = [&](int dx, int dy, int dz) {
    int x_0 = std::max(0, dx), x_1 = std::min(w, w + dx);
    int y_0 = std::max(0, dy), y_1 = std::min(h, h + dy);
    int z_0 = std::max(0, dz), z_1 = std::min(d, d + dz);
    return sl(x_0, x_1, y_0, y_1, z_0, z_1);
  };

  auto edge = [&](int dx, int dy, int dz) {
    auto s_1 = shift(dx, dy, dz);
    auto s_2 = shift(-dx, -dy, -dz);
    return torch::bitwise_and(
      torch::bitwise_not(s_1.index_into(empty)),
      s_2.index_into(empty)
    );
  };

  update(adjacencies, shift(1, 0, 0), torch::bitwise_and(edge(1, 0, 0), x_neg_mask));
  update(adjacencies, shift(-1, 0, 0), torch::bitwise_and(edge(-1, 0, 0), x_pos_mask));
  update(adjacencies, shift(0, 1, 0), torch::bitwise_and(edge(0, 1, 0), y_neg_mask));
  update(adjacencies, shift(0, -1, 0), torch::bitwise_and(edge(0, -1, 0), y_pos_mask));
  update(adjacencies, shift(0, 0, 1), torch::bitwise_and(edge(0, 0, 1), z_neg_mask));
  update(adjacencies, shift(0, 0, -1), torch::bitwise_and(edge(0, 0, -1), z_pos_mask));

  return adjacencies;
}

auto generate_mesh(const VoxelConfig& config, const VoxelTensor& tensor) {
  assert(tensor.dim() == 3);

  VoxelMesh mesh;

  // Adds 4 vertices and 2 triangles (for one cube face) to the output mesh.
  auto emit_face_geometry = [&](Vec3i origin, DirMask dir) {
    static auto positions = [] {
      using T = std::array<Vec3i, 4>;
      std::vector<T> ret;
      ret.push_back(T{{{0, 0, 0}, {0, 1, 0}, {0, 1, 1}, {0, 0, 1}}});  // X_NEG
      ret.push_back(T{{{1, 1, 0}, {1, 0, 0}, {1, 0, 1}, {1, 1, 1}}});  // X_POS
      ret.push_back(T{{{1, 0, 0}, {0, 0, 0}, {0, 0, 1}, {1, 0, 1}}});  // Y_NEG
      ret.push_back(T{{{0, 1, 0}, {1, 1, 0}, {1, 1, 1}, {0, 1, 1}}});  // Y_POS
      ret.push_back(T{{{0, 0, 0}, {1, 0, 0}, {1, 1, 0}, {0, 1, 0}}});  // Z_NEG
      ret.push_back(T{{{0, 1, 1}, {1, 1, 1}, {1, 0, 1}, {0, 0, 1}}});  // Z_POS
      return ret;
    }();

    static auto normals = [] {
      std::vector<Vec3i> ret;
      ret.push_back({-1, 0, 0});  // X_NEG
      ret.push_back({1, 0, 0});   // X_POS
      ret.push_back({0, -1, 0});  // Y_NEG
      ret.push_back({0, 1, 0});   // Y_POS
      ret.push_back({0, 0, -1});  // Y_NEG
      ret.push_back({0, 0, 1});   // Y_POS
      return ret;
    }();

    // Emit vertex indices for the two new triangles.
    auto base = mesh.vertex_count();
    mesh.triangles.push_back(base);
    mesh.triangles.push_back(base + 2);
    mesh.triangles.push_back(base + 1);
    mesh.triangles.push_back(base + 3);
    mesh.triangles.push_back(base + 2);
    mesh.triangles.push_back(base);

    // Emit vertex positions and normals for the 4 new vertices.
    auto dir_index = lg2(static_cast<int>(dir));
    for (const auto& pos_vec : positions.at(dir_index)) {
      Vec3i out = add(origin, pos_vec);
      mesh.positions.push_back(out[0]);
      mesh.positions.push_back(out[1]);
      mesh.positions.push_back(out[2]);
    }
    for (int i = 0; i < 4; i += 1) {
      Vec3i out = normals.at(dir_index);
      mesh.normals.push_back(out[0]);
      mesh.normals.push_back(out[1]);
      mesh.normals.push_back(out[2]);
    }
  };

  // Adds 4 vertex color attributes (for one cube face) to the output mesh.
  auto emit_face_colors = [&](int32_t rgba) {
    for (int i = 0; i < 4; i += 1) {
      auto [r, g, b, a] = unpack_bytes(rgba);
      mesh.colors.push_back(r);
      mesh.colors.push_back(g);
      mesh.colors.push_back(b);
      mesh.colors.push_back(a);
    }
  };

  // Create a tensor of the adjacency bitmask at each non-empty tensor cell.
  auto adjacencies = compute_adjacency_mask(tensor);

  // TODO: Pre-allocate mesh vectors to fit the exact output.

  // Iterate over each boundary voxel and emit its mesh attributes.
  auto w = tensor.size(0);
  auto h = tensor.size(1);
  auto d = tensor.size(2);
  for (int x = 0; x < w; x++) {
    for (int y = 0; y < h; y++) {
      for (int z = 0; z < d; z++) {
        Vec3i origin = {x, y, z};
        int adj = adjacencies.index({x, y, z}).item<int>();
        int key = tensor.index({x, y, z}).item<int>();

        if (adj == 0) {
          continue;
        }

        auto v = std::get<ColorVoxel>(config.defs.at(key));
        if (adj & static_cast<int>(DirMask::X_NEG)) {
          emit_face_geometry(origin, DirMask::X_NEG);
          emit_face_colors(v.rgba);
        }
        if (adj & static_cast<int>(DirMask::X_POS)) {
          emit_face_geometry(origin, DirMask::X_POS);
          emit_face_colors(v.rgba);
        }
        if (adj & static_cast<int>(DirMask::Y_NEG)) {
          emit_face_geometry(origin, DirMask::Y_NEG);
          emit_face_colors(v.rgba);
        }
        if (adj & static_cast<int>(DirMask::Y_POS)) {
          emit_face_geometry(origin, DirMask::Y_POS);
          emit_face_colors(v.rgba);
        }
        if (adj & static_cast<int>(DirMask::Z_NEG)) {
          emit_face_geometry(origin, DirMask::Z_NEG);
          emit_face_colors(v.rgba);
        }
        if (adj & static_cast<int>(DirMask::Z_POS)) {
          emit_face_geometry(origin, DirMask::Z_POS);
          emit_face_colors(v.rgba);
        }
      }
    }
  }
  return mesh;
}

} // namespace metapedia