src/Transformations/ReplaceArraysInIO/replace_dist_arrays_in_io.cpp

@@ -17,13 +17,17 @@ using std::pair;
 
 #define DEBUG_TRACE 0
 
-static bool checkDynamicArray(DIST::Array *array)
+static SgStatement* getDeclStatement(int line_num)
 {
-    for (const auto &bounds : array->GetSizes())
-        if (bounds.first == -1 || bounds.second == -1)
-            return true;
+    PTR_BFND node = current_file->firstStatement()->thebif;
+    for (; node; node = node->thread)
+    {
+        SgStatement *st = BfndMapping(node);
+        if (st->lineNumber() == line_num && isSgDeclarationStatement(st))
+            return st;
+    }
 
-    return false;
+    return NULL;
 }
 
 static SgExpression* findExprWithVariant(SgExpression* exp, int variant)
@@ -53,7 +57,6 @@ bool switchToDeclarationFile(DIST::Array* array_p,
     if (!array_p)
         return false;
 
-
     // try to find declaration from current function
     for (const auto &p : array_p->GetDeclInfo())
     {
@@ -78,6 +81,47 @@ bool switchToDeclarationFile(DIST::Array* array_p,
     return false;
 }
 
+static bool checkAssumedShape(SgStatement* decl, const string& array_name)
+{
+    SgExpression* list = decl->expr(0);
+    SgExpression* dim_list = NULL;
+
+    while (list)
+    {
+        auto *arr_ref = list->lhs();
+
+        if (arr_ref && 
+            arr_ref->symbol() && 
+            arr_ref->symbol()->identifier() && 
+            arr_ref->symbol()->identifier() == array_name)
+        {
+            dim_list = arr_ref->lhs();
+            break;
+        }
+
+        list = list->rhs();
+    }
+
+    if (!dim_list)
+    {
+        auto *dim_expr = findExprWithVariant(decl->expr(2), DIMENSION_OP);
+        if (dim_expr)
+            dim_list = dim_expr->lhs();
+    }
+
+    while (dim_list)
+    {
+        auto *lhs = dim_list->lhs();
+
+        if (lhs && lhs->variant() == DDOT && (!lhs->lhs() || !lhs->rhs()))
+            return true;
+
+        dim_list = dim_list->rhs();
+    }
+
+    return false;
+}
+
 static bool checkAssumedSize(const string &array_name, DIST::Array* array_p, const string &current_file_name, FuncInfo *current_func)
 {
     if (!array_p)
@@ -90,7 +134,10 @@ static bool checkAssumedSize(const string &array_name, DIST::Array* array_p, con
     if (!switchToDeclarationFile(array_p, decl_place, current_file_name, current_func))
         return true;
 
-    auto *st = SgStatement::getStatementByFileAndLine(decl_place.first, decl_place.second);
+    auto *st = getDeclStatement(decl_place.second);
+
+    if (!st)
+        printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
     
     SgExpression* list = st->expr(0);    
     while (list)
@@ -618,13 +665,29 @@ void replaceDistributedArraysInIO(vector<ParallelRegion*>& regions,
                                                                copied_syms);
                             
 
-                                auto* decl_stmt = SgStatement::getStatementByFileAndLine(decl_place.first, decl_place.second);
-
-                                // created declaration statement located right after original one
-                                if (!decl_stmt || !decl_stmt->lexNext())
+                                // original declaration statement
+                                auto* decl_stmt = getDeclStatement(decl_place.second);
+                                if (!decl_stmt || !isSgDeclarationStatement(decl_stmt))
                                     printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
 
-                                decl_stmt = decl_stmt->lexNext();
+                                auto dynamic_array = checkAssumedShape(decl_stmt, array_name);  
+
+                                while (decl_stmt && !isSgExecutableStatement(decl_stmt))
+                                {
+                                    if (isSgDeclarationStatement(decl_stmt) &&
+                                        decl_stmt->expr(0) &&
+                                        decl_stmt->expr(0)->lhs() &&
+                                        decl_stmt->expr(0)->lhs()->symbol() == copied_symbol.second)
+                                    {
+                                        break;
+                                    }
+
+                                    decl_stmt = decl_stmt->lexNext();
+                                }
+                                
+                                // created declaration statement
+                                if (!decl_stmt || !isSgDeclarationStatement(decl_stmt))
+                                    printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
 
                                 // insert !$SPF ANALYSIS(PROCESS_PRIVATE(array)) directive before declaration statement
                                 
@@ -641,8 +704,8 @@ void replaceDistributedArraysInIO(vector<ParallelRegion*>& regions,
                             
                                 created_copies.insert({ array_to_copy, copied_symbol.second });
 
-                                // make array copy allocatable in case of main array shape not constant
-                                if(checkDynamicArray(array_p))
+                                // make array copy allocatable in case of assumed-shape array
+                                if(dynamic_array)
                                 {
                                     // insert allocatable keyword in declaration
                                     auto *kword_list = decl_stmt->expr(2);