OpenSCAD Tree Flatterning: push transforms down and bubble unions up

RepeatedNodesDetector.h

#pragma once

#include <unordered_set>
#include <unordered_map>
// #include <flat_set>
#include "node.h"

class RepeatedNodesDetector : public NodeVisitor
{
public:
  typedef std::unordered_map<std::string, std::unordered_set<const AbstractNode*>> Occurrences;

  RepeatedNodesDetector(const Tree& tree, Occurrences& key_occurrences) : tree(tree), key_occurrences(key_occurrences) {}

  Response visit(State& state, const AbstractNode& node) override {
    if (state.isPrefix()) {
      auto key = this->tree.getIdString(node);
      key_occurrences[key].insert(&node);
    }
    return Response::ContinueTraversal;
  }

  std::unordered_set<const AbstractNode*> getRepeatedNodes() const {
    std::unordered_set<const AbstractNode*> nodes;
    for (auto &pair : key_occurrences) {
      if (pair.second.size() > 1) {
        for (auto node : pair.second) {
          nodes.insert(node);
        }
      }
    }
    return std::move(nodes);
  }
  std::unordered_set<std::string> getRepeatedKeys() const {
    std::unordered_set<std::string> keys;
    for (auto &pair : key_occurrences) {
      if (pair.second.size() > 1) {
        keys.insert(pair.first);
      }
    }
    return std::move(keys);
  }

private:
  const Tree& tree;
  // std::unordered_map<std::string, size_t>& key_occurrences;
  Occurrences& key_occurrences;
};

TreeFlattener.h

#pragma once

#include <unordered_set>
#include <unordered_map>
#include <stack>
#include "node.h"
#include "TransformNode.h"
#include "CsgOpNode.h"
#include "ColorNode.h"
#include "RepeatedNodesDetector.h"

/**
 * Flattens unions of unions, intersections of intersections, and pushes down colors and transforms as far as it can.
 *
 * This creates flatter trees, with higher arity in commutative operations and more lazy unions at the top.
 * 
 * Skips special nodes and subtrees which are repeated.
*/
class TreeFlattener
{
  struct State {
    std::optional<Transform3d> transform;
    std::optional<Color4f> color;
  };

  const Tree& tree;
  std::unordered_set<const AbstractNode*> repeatedNodes;
  bool traverseHulls;

public:
  TreeFlattener(const Tree& tree)
    : tree(tree), traverseHulls(false)
  {
    RepeatedNodesDetector::Occurrences key_occurrences;
    RepeatedNodesDetector detector(tree, key_occurrences);
    repeatedNodes = detector.getRepeatedNodes();
  }

  void setTraverseHulls(bool value) {
    traverseHulls = value;
  }
  
  shared_ptr<const AbstractNode> flatten(const shared_ptr<const AbstractNode>& node, const State &state = State {}, bool allowRecursion = true) {
    if (!node) return node;

    using Children = std::vector<shared_ptr<const AbstractNode>>;

    if (canInline(node)) {
      auto makeOp = [&](OpenSCADOperator op, const Children& children, bool forceFlattenIfNewNode = false) -> shared_ptr<const AbstractNode> {
        auto needsFlatten = allowRecursion && (node->children != children) || forceFlattenIfNewNode;
        
        if (children.size() == 1) {
          return children[0];
        }

        shared_ptr<AbstractNode> ret;
        if (op == OpenSCADOperator::UNION) {
          ret = lazyUnionNode(node->modinst);
        } else {
          ret = make_shared<CsgOpNode>(node->modinst, op);
        }
        ret->children = children;

        if (needsFlatten) {
          return flatten(ret, State {}, /* allowRecursion= */ false);
        } else {
          return ret;
        }
      };

      if (dynamic_pointer_cast<const ListNode>(node) || 
          dynamic_pointer_cast<const GroupNode>(node)) {
        Children children;
        flattenChildren(node, children, OpenSCADOperator::UNION, state, allowRecursion);
        return makeOp(OpenSCADOperator::UNION, children);
      } else if (auto csgOpNode = dynamic_pointer_cast<const CsgOpNode>(node)) {
        if (csgOpNode->type == OpenSCADOperator::UNION ||
            csgOpNode->type == OpenSCADOperator::INTERSECTION)
        {
          Children children;
          flattenChildren(csgOpNode, children, csgOpNode->type, state, allowRecursion);

          return makeOp(csgOpNode->type, children);
        }
      } else if (auto transformNode = dynamic_pointer_cast<const TransformNode>(node)) {
        Children children;
        flattenChildren(transformNode, children, OpenSCADOperator::UNION, State {}, allowRecursion);
        if (children.size() > 1 ||
            children.size() == 1 && dynamic_pointer_cast<const TransformNode>(children[0])) {
          State newState = state;
          if (auto transform = state.transform) {
            newState.transform = *transform * transformNode->matrix;
          } else {
            newState.transform = transformNode->matrix;
          }

          for (auto &child : children) {
            child = flatten(child, newState);
          }

          return makeOp(OpenSCADOperator::UNION, children, /* forceFlattenIfNewNode= */ true);
        }
      } else if (auto colorNode = dynamic_pointer_cast<const ColorNode>(node)) {
        Children children;
        flattenChildren(colorNode, children, OpenSCADOperator::UNION, State {}, allowRecursion);
        if (children.size() > 1 ||
            children.size() == 1 && dynamic_pointer_cast<const ColorNode>(children[0])) {
          State newState = state;
          if (auto color = state.color) {
            newState.color = *color;
          }

          for (auto &child : children) {
            child = flatten(child, newState);
          }

          return makeOp(OpenSCADOperator::UNION, children, /* forceFlattenIfNewNode= */ true);
        }
      }
    }
    
    // if (canTraverse(node)) {
    //   return transformChildren(node, state);
    // }
    
    // if (auto leafNode = dynamic_pointer_cast<LeafNode>(node)) {
    //   return wrapWithState(leafNode);
    // }
    return wrapWithState(node, state);
  }

private:

  // flattenAs<ListNode>(List(a, List(b, c, Union(d, e)))), Union) -> 
  void flattenChildren(const shared_ptr<const AbstractNode>& node, std::vector<shared_ptr<const AbstractNode>> &out, const std::optional<OpenSCADOperator> allowedOp, const State &state, bool allowRecursion) {
    if (!node) {
      return;
    }
    for (auto child : node->children) {
      if (canInline(child)) {
        if (dynamic_pointer_cast<const ListNode>(child) || dynamic_pointer_cast<const GroupNode>(child)) {
          flattenChildren(child, out, allowedOp, state, allowRecursion);
          continue;
        }
        if (auto csgNode = dynamic_pointer_cast<const CsgOpNode>(child)) {
          if (allowedOp && csgNode->type == *allowedOp && isAssociative(*allowedOp)) {
            flattenChildren(csgNode, out, allowedOp, state, allowRecursion);
            continue;
          }
        }
      }
      if (allowRecursion) {
        out.push_back(flatten(child, state));
      } else {
        out.push_back(wrapWithState(child, state));
      }
    }
  }

  bool isAssociative(OpenSCADOperator op) const {
    return op == OpenSCADOperator::UNION ||
      op == OpenSCADOperator::INTERSECTION ||
      op == OpenSCADOperator::HULL;
  }

  bool hasSpecialTags(const shared_ptr<const AbstractNode>& node) const {
    if (!node || !node->modinst) {
      return false;
    }
    
    // Not testing for isRoot, as we'll be called on the actual tree root to only render what's below it.
    return node->modinst->isBackground() ||
        node->modinst->isHighlight();
  }

  shared_ptr<const AbstractNode> wrapWithState(const shared_ptr<const AbstractNode>& node, const State &state) const {
    auto res = node;
    if (auto transform = state.transform) {
      auto transformNode = make_shared<TransformNode>(node->modinst, "??");
      transformNode->children.push_back(res);
      transformNode->matrix = *transform;
      res = transformNode;
    }
    if (auto color = state.color) {
      auto colorNode = make_shared<ColorNode>(node->modinst);
      colorNode->children.push_back(res);
      colorNode->color = *color;
      res = colorNode;
    }
    return res;
  }

  // shared_ptr<const AbstractNode> transformChildren(const shared_ptr<const AbstractNode>& node, const State &state) const {
  //   assert(false && "not implemented");
  //   return node; 
  // }

  bool canTraverse(const shared_ptr<const AbstractNode>& node) const {
    if (hasSpecialTags(node)) {
      return false;
    }
    if (auto csgOpNode = dynamic_pointer_cast<const CsgOpNode>(node)) {
      switch (csgOpNode->type) {
        case OpenSCADOperator::UNION:
        case OpenSCADOperator::INTERSECTION:
        case OpenSCADOperator::DIFFERENCE:
          return true;
        case OpenSCADOperator::HULL:
          return traverseHulls;
        default:
          return false;
      }
    }
    return
        dynamic_pointer_cast<const ListNode>(node) ||
        dynamic_pointer_cast<const GroupNode>(node);
  }
  
  bool canInline(const shared_ptr<const AbstractNode>& node) const {
    if (!node || hasSpecialTags(node)) {
      return false;
    }
    if (repeatedNodes.find(node.get()) != repeatedNodes.end()) {
      return false;
    }
    return true;
  }
};