Simplify UnitProto after building a CsgUnit (#1415)

src/orange/orangeinp/CsgTreeUtils.cc

@@ -173,7 +173,7 @@ void simplify(CsgTree* tree, NodeId start)
  * This is required if the tree's logic expression is used with
  * \c InfixEvaluator as negated joins are not supported.
  */
-CsgTree transform_negated_joins(CsgTree const& tree)
+SimplifiedCsgTree transform_negated_joins(CsgTree const& tree)
 {
     return detail::DeMorganSimplifier{tree}();
 }

src/orange/orangeinp/CsgTreeUtils.hh

@@ -12,13 +12,27 @@
 
 #include "orange/OrangeTypes.hh"
 
+#include "CsgTree.hh"
 #include "CsgTypes.hh"
 
 namespace celeritas
 {
 namespace orangeinp
 {
-class CsgTree;
+
+//---------------------------------------------------------------------------//
+/*!
+ * Result of a DeMorgan simplification.
+ */
+struct SimplifiedCsgTree
+{
+    //! The simplified tree
+    CsgTree tree;
+    //! Has the same size as the original tree, indexed by old node id, with
+    //! the value pointing to the equivalent node in the simplified tree
+    std::vector<NodeId> new_nodes;
+};
+
 //---------------------------------------------------------------------------//
 
 // Replace a node in the tree with a boolean constant
@@ -33,7 +47,7 @@ orangeinp::NodeId simplify_up(CsgTree* tree, orangeinp::NodeId start);
 void simplify(CsgTree* tree, orangeinp::NodeId start);
 
 // Replace ~&(xs...) with |(~xs...) and ~|(xs...) with &(~xs...)
-[[nodiscard]] CsgTree transform_negated_joins(CsgTree const& tree);
+[[nodiscard]] SimplifiedCsgTree transform_negated_joins(CsgTree const& tree);
 
 // Transform a CSG node into a string expression
 [[nodiscard]] std::string

src/orange/orangeinp/CsgTypes.hh

@@ -71,6 +71,16 @@ struct Joined
 //! Generic node
 using Node = std::variant<True, False, Aliased, Negated, Surface, Joined>;
 
+/*!
+ * Optional transformations to apply when building a CsgUnit.
+ */
+enum class UnitSimplification : size_type
+{
+    none = 0,  //!< No simplification
+    infix_logic,  //!< CsgTree suitable for infix logic evaluation
+    size_
+};
+
 //---------------------------------------------------------------------------//
 // Equality operators
 //---------------------------------------------------------------------------//

src/orange/orangeinp/UnitProto.cc

@@ -458,6 +458,11 @@ auto UnitProto::build(Tol const& tol, BBox const& bbox) const -> Unit
                                write_node_labels);
         }
 
+        if (input_.simplification == UnitSimplification::infix_logic)
+        {
+            unit_builder.simplifiy_joins();
+        }
+
         /*! \todo We can sometimes eliminate CSG surfaces and nodes that aren't
          * used by the actual volumes>
          */

src/orange/orangeinp/UnitProto.hh

@@ -12,6 +12,7 @@
 
 #include "geocel/Types.hh"
 #include "orange/OrangeTypes.hh"
+#include "orange/orangeinp/CsgTypes.hh"
 #include "orange/transform/VariantTransform.hh"
 
 #include "ProtoInterface.hh"
@@ -130,6 +131,7 @@ class UnitProto : public ProtoInterface
         std::vector<DaughterInput> daughters;
         BoundaryInput boundary;
         std::string label;
+        UnitSimplification simplification{UnitSimplification::none};
 
         // True if fully defined
         explicit inline operator bool() const;

src/orange/orangeinp/detail/CsgUnitBuilder.cc

@@ -11,6 +11,7 @@
 #include "corecel/io/Logger.hh"
 #include "corecel/io/StreamableVariant.hh"
 #include "orange/OrangeData.hh"
+#include "orange/orangeinp/CsgTreeUtils.hh"
 #include "orange/transform/TransformIO.hh"
 #include "orange/transform/TransformSimplifier.hh"
 
@@ -193,6 +194,48 @@ void CsgUnitBuilder::fill_volume(LocalVolumeId v,
     CELER_ENSURE(is_filled(unit_->fills[v.unchecked_get()]));
 }
 
+//---------------------------------------------------------------------------//
+/*!
+ * Simplify negated joins for Infix evaluation.
+ *
+ * Apply DeMorgan simplification to use the \c CsgUnit in infix evaluation.
+ * \c NodeId indexing in the \c CsgTree are invalidated after calling this,
+ * \c CsgUnit data is updated to point to the simplified tree \c NodeId but any
+ * previously cached \c NodeId is invalid.
+ */
+void CsgUnitBuilder::simplifiy_joins()
+{
+    auto& tree = unit_->tree;
+    auto simplification = transform_negated_joins(tree);
+    CELER_ASSERT(tree.size() == simplification.new_nodes.size());
+    std::vector<std::set<CsgUnit::Metadata>> md;
+    md.resize(simplification.tree.size());
+
+    std::map<NodeId, CsgUnit::Region> regions;
+
+    for (auto node_id : range(tree.size()))
+    {
+        if (auto equivalent_node = simplification.new_nodes[node_id])
+        {
+            CELER_EXPECT(equivalent_node < md.size());
+            md[equivalent_node.unchecked_get()]
+                = std::move(unit_->metadata[node_id]);
+            regions[equivalent_node]
+                = std::move(unit_->regions[NodeId{node_id}]);
+        }
+        else if (unit_->regions.find(NodeId{node_id}) != unit_->regions.end()
+                 || !unit_->metadata[node_id].empty())
+        {
+            CELER_LOG(warning)
+                << "While simplifying node '" << node_id
+                << "': has metadata or region but no equivalent node";
+        }
+    }
+    unit_->metadata = std::move(md);
+    unit_->regions = std::move(regions);
+    unit_->tree = std::move(simplification.tree);
+}
+
 //---------------------------------------------------------------------------//
 /*!
  * Get a variant surface from a node ID.

src/orange/orangeinp/detail/CsgUnitBuilder.hh

@@ -107,6 +107,9 @@ class CsgUnitBuilder
     void
     fill_volume(LocalVolumeId, UniverseId, VariantTransform const& transform);
 
+    // Simplify negated joins for Infix evaluation
+    void simplifiy_joins();
+
   private:
     CsgUnit* unit_;
     Tol tol_;

src/orange/orangeinp/detail/DeMorganSimplifier.cc

@@ -16,6 +16,7 @@
 #include "corecel/cont/Range.hh"
 #include "orange/OrangeTypes.hh"
 #include "orange/orangeinp/CsgTree.hh"
+#include "orange/orangeinp/CsgTreeUtils.hh"
 #include "orange/orangeinp/CsgTypes.hh"
 
 namespace celeritas
@@ -87,10 +88,18 @@ DeMorganSimplifier::DeMorganSimplifier(CsgTree const& tree)
  * Perform the simplification. The state of the instance isn't cleared, so only
  * call this once.
  */
-CsgTree DeMorganSimplifier::operator()()
+SimplifiedCsgTree DeMorganSimplifier::operator()()
 {
     this->find_join_negations();
-    return this->build_simplified_tree();
+    auto simplified_tree{this->build_simplified_tree()};
+    std::vector<NodeId> equivalent_nodes;
+    equivalent_nodes.reserve(tree_.size());
+    for (auto node_id : range(tree_.size()))
+    {
+        equivalent_nodes.push_back(
+            node_ids_translation_[node_id].equivalent_node());
+    }
+    return {simplified_tree, equivalent_nodes};
 }
 
 //---------------------------------------------------------------------------//

src/orange/orangeinp/detail/DeMorganSimplifier.hh

@@ -17,6 +17,7 @@ namespace celeritas
 {
 namespace orangeinp
 {
+struct SimplifiedCsgTree;
 namespace detail
 {
 //---------------------------------------------------------------------------//
@@ -46,7 +47,7 @@ class DeMorganSimplifier
     explicit DeMorganSimplifier(CsgTree const&);
 
     // Perform the simplification
-    CsgTree operator()();
+    SimplifiedCsgTree operator()();
 
   private:
     //! CsgTree node 0 is always True{} and can't be the parent of any node