From 8a297abd0abffb6af7c8cbd4ffcd5e5f554c8ef0 Mon Sep 17 00:00:00 2001
From: James Wright <james@jameswright.xyz>
Date: Thu, 10 Oct 2024 21:32:39 -0600
Subject: [PATCH] feat(op): Add CeedCompositeOperatorGetSubByName()

Co-authored-by: Jeremy L Thompson <jeremy@jeremylt.org>
---
 include/ceed/ceed.h       |  1 +
 interface/ceed-operator.c | 34 ++++++++++++++++++++++++++++++++++
 tests/t520-operator.c     | 12 ++++++++++++
 3 files changed, 47 insertions(+)

diff --git a/include/ceed/ceed.h b/include/ceed/ceed.h
index d847525da8..1660d724a9 100644
--- a/include/ceed/ceed.h
+++ b/include/ceed/ceed.h
@@ -416,6 +416,7 @@ CEED_EXTERN int  CeedOperatorIsAtPoints(CeedOperator op, bool *is_at_points);
 CEED_EXTERN int  CeedCompositeOperatorAddSub(CeedOperator composite_op, CeedOperator sub_op);
 CEED_EXTERN int  CeedCompositeOperatorGetNumSub(CeedOperator op, CeedInt *num_suboperators);
 CEED_EXTERN int  CeedCompositeOperatorGetSubList(CeedOperator op, CeedOperator **sub_operators);
+CEED_EXTERN int  CeedCompositeOperatorGetSubByName(CeedOperator op, const char *op_name, CeedOperator *sub_op);
 CEED_EXTERN int  CeedOperatorCheckReady(CeedOperator op);
 CEED_EXTERN int  CeedOperatorGetActiveVectorLengths(CeedOperator op, CeedSize *input_size, CeedSize *output_size);
 CEED_EXTERN int  CeedOperatorSetQFunctionAssemblyReuse(CeedOperator op, bool reuse_assembly_data);
diff --git a/interface/ceed-operator.c b/interface/ceed-operator.c
index 5881846665..065c5a2cda 100644
--- a/interface/ceed-operator.c
+++ b/interface/ceed-operator.c
@@ -1284,6 +1284,40 @@ int CeedCompositeOperatorGetSubList(CeedOperator op, CeedOperator **sub_operator
   return CEED_ERROR_SUCCESS;
 }
 
+/**
+  @brief Get a sub `CeedOperator` of a composite `CeedOperator` from its name.
+
+  `sub_op` is set to `NULL` if the sub operator is not found.
+
+  Note: Calling this function asserts that setup is complete and sets the `CeedOperator` as immutable.
+
+  @param[in]  op      Composite `CeedOperator`
+  @param[in]  op_name Name of desired sub `CeedOperator`
+  @param[out] sub_op  Sub `CeedOperator` corresponding to the name
+
+  @return An error code: 0 - success, otherwise - failure
+
+  @ref Advanced
+**/
+int CeedCompositeOperatorGetSubByName(CeedOperator op, const char *op_name, CeedOperator *sub_op) {
+  bool          is_composite;
+  CeedInt       num_sub_ops;
+  CeedOperator *sub_ops;
+
+  CeedCall(CeedOperatorIsComposite(op, &is_composite));
+  CeedCheck(is_composite, CeedOperatorReturnCeed(op), CEED_ERROR_MINOR, "Only defined for a composite operator");
+  *sub_op = NULL;
+  CeedCall(CeedCompositeOperatorGetNumSub(op, &num_sub_ops));
+  CeedCall(CeedCompositeOperatorGetSubList(op, &sub_ops));
+  for (CeedInt i = 0; i < num_sub_ops; i++) {
+    if (sub_ops[i]->name && !strcmp(op_name, sub_ops[i]->name)) {
+      *sub_op = sub_ops[i];
+      return CEED_ERROR_SUCCESS;
+    }
+  }
+  return CEED_ERROR_SUCCESS;
+}
+
 /**
   @brief Check if a `CeedOperator` is ready to be used.
 
diff --git a/tests/t520-operator.c b/tests/t520-operator.c
index 9035234ffa..31fceb9c97 100644
--- a/tests/t520-operator.c
+++ b/tests/t520-operator.c
@@ -111,6 +111,7 @@ int main(int argc, char **argv) {
   CeedOperatorSetField(op_mass_tet, "rho", elem_restriction_q_data_tet, CEED_BASIS_NONE, q_data_tet);
   CeedOperatorSetField(op_mass_tet, "u", elem_restriction_u_tet, basis_u_tet, CEED_VECTOR_ACTIVE);
   CeedOperatorSetField(op_mass_tet, "v", elem_restriction_u_tet, basis_u_tet, CEED_VECTOR_ACTIVE);
+  CeedOperatorSetName(op_mass_tet, "mass tet");
 
   // Set up Hex Elements
   // -- Restrictions
@@ -154,6 +155,7 @@ int main(int argc, char **argv) {
   CeedOperatorSetField(op_mass_hex, "rho", elem_restriction_q_data_hex, CEED_BASIS_NONE, q_data_hex);
   CeedOperatorSetField(op_mass_hex, "u", elem_restriction_u_hex, basis_u_hex, CEED_VECTOR_ACTIVE);
   CeedOperatorSetField(op_mass_hex, "v", elem_restriction_u_hex, basis_u_hex, CEED_VECTOR_ACTIVE);
+  CeedOperatorSetName(op_mass_hex, "mass hex");
 
   // Set up Composite Operators
   // -- Create
@@ -168,6 +170,16 @@ int main(int argc, char **argv) {
   CeedCompositeOperatorAddSub(op_mass, op_mass_tet);
   CeedCompositeOperatorAddSub(op_mass, op_mass_hex);
 
+  {  // Test CeedCompositeOperatorGetSubByName
+    CeedOperator op_byname;
+
+    CeedCompositeOperatorGetSubByName(op_mass, "mass hex", &op_byname);
+    if (op_byname != op_mass_hex) printf("CeedCompositeOperatorGetSubByName returned incorrect Sub Operator");
+
+    CeedCompositeOperatorGetSubByName(op_mass, "asdf", &op_byname);
+    if (op_byname != NULL) printf("CeedCompositeOperatorGetSubByName returned non-NULL for non-existent Sub Operator");
+  }
+
   // Apply Setup Operator
   CeedOperatorApply(op_setup, x, CEED_VECTOR_NONE, CEED_REQUEST_IMMEDIATE);