Implement getting associative domain type from subs

Signed-off-by: Anna Rift <[email protected]>
chapel-lang · Nov 25, 2024 · 9b06d78 · 9b06d78
1 parent 53a38b2
commit 9b06d78
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 9 deletions.
diff --git a/frontend/include/chpl/types/DomainType.h b/frontend/include/chpl/types/DomainType.h
@@ -98,8 +98,9 @@ class DomainType final : public CompositeType {
 
   /** Return an associative domain type */
   static const DomainType* getAssociativeType(Context* context,
-                                              const QualifiedType& idxType,
-                                              const QualifiedType& parSafe);
+                                 const QualifiedType& instance,
+                                 const QualifiedType& idxType,
+                                 const QualifiedType& parSafe);
 
   /** Get the default distribution type */
   static const QualifiedType& getDefaultDistType(Context* context);

diff --git a/frontend/lib/resolution/InitResolver.cpp b/frontend/lib/resolution/InitResolver.cpp
@@ -295,9 +295,9 @@ static const DomainType* domainTypeFromSubsHelper(
       if (auto instanceBct = instanceCt->basicClassType()) {
         // Get BaseRectangularDom parent subs for rectangular domain info
         if (auto baseDom = instanceBct->parentClassType()) {
-          auto& rf = fieldsForTypeDecl(context, baseDom,
-                                       DefaultsPolicy::IGNORE_DEFAULTS);
           if (baseDom->id().symbolPath() == "ChapelDistribution.BaseRectangularDom") {
+            auto& rf = fieldsForTypeDecl(context, baseDom,
+                                         DefaultsPolicy::IGNORE_DEFAULTS);
             CHPL_ASSERT(rf.numFields() == 3);
             QualifiedType rank;
             QualifiedType idxType;
@@ -315,7 +315,24 @@ static const DomainType* domainTypeFromSubsHelper(
             return DomainType::getRectangularType(context, instanceQt, rank,
                                                   idxType, strides);
           } else if (baseDom->id().symbolPath() == "ChapelDistribution.BaseAssociativeDom") {
-            // TODO: support associative domains
+            // Currently the relevant associative domain fields are defined
+            // on all the children of BaseAssociativeDom, so get information
+            // from there.
+            auto& rf = fieldsForTypeDecl(context, instanceBct,
+                                         DefaultsPolicy::IGNORE_DEFAULTS);
+            CHPL_ASSERT(rf.numFields() >= 2);
+            QualifiedType idxType;
+            QualifiedType parSafe;
+            for (int i = 0; i < rf.numFields(); i++) {
+              if (rf.fieldName(i) == "idxType") {
+                idxType = rf.fieldType(i);
+              } else if (rf.fieldName(i) == "parSafe") {
+                parSafe = rf.fieldType(i);
+              }
+            }
+
+            return DomainType::getAssociativeType(context, instanceQt, idxType,
+                                                  parSafe);
           } else if (baseDom->id().symbolPath() == "ChapelDistribution.BaseSparseDom") {
             // TODO: support sparse domains
           } else {

diff --git a/frontend/lib/resolution/prims.cpp b/frontend/lib/resolution/prims.cpp
@@ -293,6 +293,7 @@ static QualifiedType computeDomainType(Context* context, const CallInfo& ci) {
     return QualifiedType(QualifiedType::TYPE, type);
   } else if (ci.numActuals() == 2) {
     auto type = DomainType::getAssociativeType(context,
+                                               QualifiedType(),
                                                ci.actual(0).type(),
                                                ci.actual(1).type());
     return QualifiedType(QualifiedType::TYPE, type);

diff --git a/frontend/lib/types/DomainType.cpp b/frontend/lib/types/DomainType.cpp
@@ -118,17 +118,35 @@ DomainType::getRectangularType(Context* context,
 
 const DomainType*
 DomainType::getAssociativeType(Context* context,
+                               const QualifiedType& instance,
                                const QualifiedType& idxType,
                                const QualifiedType& parSafe) {
+  auto genericDomain = getGenericDomainType(context);
+
   SubstitutionsMap subs;
-  // TODO: assert validity of sub types
   subs.emplace(ID(UniqueString(), 0, 0), idxType);
+  CHPL_ASSERT(idxType.isType());
   subs.emplace(ID(UniqueString(), 1, 0), parSafe);
+  CHPL_ASSERT(parSafe.isParam() && parSafe.param() &&
+              parSafe.param()->isBoolParam());
+
+  // Add substitution for _instance field
+  auto& rf = fieldsForTypeDecl(context, genericDomain,
+                               resolution::DefaultsPolicy::IGNORE_DEFAULTS,
+                               /* syntaxOnly */ true);
+  ID instanceFieldId;
+  for (int i = 0; i < rf.numFields(); i++) {
+    if (rf.fieldName(i) == USTR("_instance")) {
+      instanceFieldId = rf.fieldDeclId(i);
+      break;
+    }
+  }
+  subs.emplace(instanceFieldId, instance);
+
   auto name = UniqueString::get(context, "_domain");
   auto id = getDomainID(context);
-  auto instantiatedFrom = getGenericDomainType(context);
-  return getDomainType(context, id, name, instantiatedFrom, subs,
-                       DomainType::Kind::Associative).get();
+  return getDomainType(context, id, name, /* instantiatedFrom */ genericDomain,
+                       subs, DomainType::Kind::Associative).get();
 }
 
 const QualifiedType& DomainType::getDefaultDistType(Context* context) {