From 8b206234f3c1e9fd32490e5abe8a3aac344c93af Mon Sep 17 00:00:00 2001
From: Steven Perron <stevenperron@google.com>
Date: Mon, 12 Feb 2024 14:52:41 -0500
Subject: [PATCH] Refactor code to get the number of members.

---
 source/opt/copy_prop_arrays.cpp | 48 ++++++++++++++++++---------------
 1 file changed, 27 insertions(+), 21 deletions(-)

diff --git a/source/opt/copy_prop_arrays.cpp b/source/opt/copy_prop_arrays.cpp
index 322d7123fef..398839818d0 100644
--- a/source/opt/copy_prop_arrays.cpp
+++ b/source/opt/copy_prop_arrays.cpp
@@ -35,6 +35,32 @@ bool IsDebugDeclareOrValue(Instruction* di) {
          dbg_opcode == CommonDebugInfoDebugValue;
 }
 
+// Returns the number of members in |type|.  If |type| is not a composite type
+// or the number of components is not known at compile time, the return value
+// will be 0.
+uint32_t GetNumberOfMembers(const analysis::Type* type, IRContext* context) {
+  if (const analysis::Struct* struct_type = type->AsStruct()) {
+    return static_cast<uint32_t>(struct_type->element_types().size());
+  } else if (const analysis::Array* array_type = type->AsArray()) {
+    const analysis::Constant* length_const =
+        context->get_constant_mgr()->FindDeclaredConstant(
+            array_type->LengthId());
+
+    if (length_const == nullptr) {
+      // This can happen if the length is an OpSpecConstant.
+      return 0;
+    }
+    assert(length_const->type()->AsInteger());
+    return length_const->GetU32();
+  } else if (const analysis::Vector* vector_type = type->AsVector()) {
+    return vector_type->element_count();
+  } else if (const analysis::Matrix* matrix_type = type->AsMatrix()) {
+    return matrix_type->element_count();
+  } else {
+    return 0;
+  }
+}
+
 }  // namespace
 
 Pass::Status CopyPropagateArrays::Process() {
@@ -808,28 +834,8 @@ uint32_t CopyPropagateArrays::MemoryObject::GetNumberOfMembers() {
   std::vector<uint32_t> access_indices = GetAccessIds();
   type = type_mgr->GetMemberType(type, access_indices);
 
-  if (const analysis::Struct* struct_type = type->AsStruct()) {
-    return static_cast<uint32_t>(struct_type->element_types().size());
-  } else if (const analysis::Array* array_type = type->AsArray()) {
-    const analysis::Constant* length_const =
-        context->get_constant_mgr()->FindDeclaredConstant(
-            array_type->LengthId());
-
-    if (length_const == nullptr) {
-      // This can happen if the length is an OpSpecConstant.
-      return 0;
-    }
-    assert(length_const->type()->AsInteger());
-    return length_const->GetU32();
-  } else if (const analysis::Vector* vector_type = type->AsVector()) {
-    return vector_type->element_count();
-  } else if (const analysis::Matrix* matrix_type = type->AsMatrix()) {
-    return matrix_type->element_count();
-  } else {
-    return 0;
-  }
+  return opt::GetNumberOfMembers(type, context);
 }
-
 template <class iterator>
 CopyPropagateArrays::MemoryObject::MemoryObject(Instruction* var_inst,
                                                 iterator begin, iterator end)