camfort · raehik · Aug 18, 2021
diff --git a/src/Language/Fortran/AST.hs b/src/Language/Fortran/AST.hs
@@ -86,6 +86,7 @@ module Language.Fortran.AST
   , DoSpecification(..)
   , ProgramUnitName(..)
   , Kind
+  , BlockConstructStart(..)
 
   -- * Node annotations & related typeclasses
   , A0
@@ -301,6 +302,17 @@ programUnitSubprograms PUComment{}                 = Nothing
 newtype Comment a = Comment String
   deriving (Eq, Show, Data, Typeable, Generic, Functor)
 
+-- | Common data related to the start of block constructs.
+data BlockConstructStart a =
+    BlockConstructStart a
+                        SrcSpan
+                        -- ^ original block start statement 'SrcSpan'
+                        (Maybe (Expression a))
+                        -- ^ label
+                        (Maybe String)
+                        -- ^ name
+  deriving (Eq, Show, Data, Typeable, Generic, Functor)
+
 data Block a =
     BlStatement a SrcSpan
                 (Maybe (Expression a))       -- ^ Label
@@ -329,16 +341,14 @@ data Block a =
                 (Maybe (Expression a))       -- ^ Label to END SELECT
 
   | BlDo        a SrcSpan
-                (Maybe (Expression a))       -- ^ Label
-                (Maybe String)               -- ^ Construct name
+                (BlockConstructStart a)
                 (Maybe (Expression a))       -- ^ Target label
                 (Maybe (DoSpecification a))  -- ^ Do Specification
                 [ Block a ]                  -- ^ Body
                 (Maybe (Expression a))       -- ^ Label to END DO
 
   | BlDoWhile   a SrcSpan
-                (Maybe (Expression a))       -- ^ Label
-                (Maybe String)               -- ^ Construct name
+                (BlockConstructStart a)
                 (Maybe (Expression a))       -- ^ Target label
                 (Expression a)               -- ^ Condition
                 [ Block a ]                  -- ^ Body
@@ -754,6 +764,7 @@ instance FirstParameter (Declarator a) a
 instance FirstParameter (DimensionDeclarator a) a
 instance FirstParameter (ControlPair a) a
 instance FirstParameter (AllocOpt a) a
+instance FirstParameter (BlockConstructStart a) a
 
 instance SecondParameter (ProgramUnit a) SrcSpan
 instance SecondParameter (Prefix a) SrcSpan
@@ -783,6 +794,7 @@ instance SecondParameter (Declarator a) SrcSpan
 instance SecondParameter (DimensionDeclarator a) SrcSpan
 instance SecondParameter (ControlPair a) SrcSpan
 instance SecondParameter (AllocOpt a) SrcSpan
+instance SecondParameter (BlockConstructStart a) SrcSpan
 
 instance Annotated (AList t)
 instance Annotated ProgramUnit
@@ -811,6 +823,7 @@ instance Annotated Declarator
 instance Annotated DimensionDeclarator
 instance Annotated ControlPair
 instance Annotated AllocOpt
+instance Annotated BlockConstructStart
 
 instance Spanned (ProgramUnit a)
 instance Spanned (Prefix a)
@@ -840,6 +853,7 @@ instance Spanned (Declarator a)
 instance Spanned (DimensionDeclarator a)
 instance Spanned (ControlPair a)
 instance Spanned (AllocOpt a)
+instance Spanned (BlockConstructStart a)
 
 instance Spanned (ProgramFile a) where
   getSpan (ProgramFile _ pus) =
@@ -854,25 +868,30 @@ class Labeled f where
   getLastLabel :: f a -> Maybe (Expression a)
   setLabel :: f a -> Expression a -> f a
 
+instance Labeled BlockConstructStart where
+    getLabel (BlockConstructStart _ _ l _) = l
+    getLastLabel = const Nothing
+    setLabel (BlockConstructStart a ss _ s) l = BlockConstructStart a ss (Just l) s
+
 instance Labeled Block where
   getLabel (BlStatement _ _ l _) = l
   getLabel (BlIf _ _ l _ _ _ _) = l
   getLabel (BlCase _ _ l _ _ _ _ _) = l
-  getLabel (BlDo _ _ l _ _ _ _ _) = l
-  getLabel (BlDoWhile _ _ l _ _ _ _ _) = l
+  getLabel (BlDo _ _ x _ _ _ _) = getLabel x
+  getLabel (BlDoWhile _ _ x _ _ _ _) = getLabel x
   getLabel _ = Nothing
 
   getLastLabel b@BlStatement{} = getLabel b
   getLastLabel (BlIf _ _ _ _ _ _ l) = l
   getLastLabel (BlCase _ _ _ _ _ _ _ l) = l
-  getLastLabel (BlDo _ _ _ _ _ _ _ l) = l
-  getLastLabel (BlDoWhile _ _ _ _ _ _ _ l) = l
+  getLastLabel (BlDo _ _ _ _ _ _ l) = l
+  getLastLabel (BlDoWhile _ _ _ _ _ _ l) = l
   getLastLabel _ = Nothing
 
   setLabel (BlStatement a s _ st) l = BlStatement a s (Just l) st
   setLabel (BlIf a s _ mn conds bs el) l = BlIf a s (Just l) mn conds bs el
-  setLabel (BlDo a s _ mn tl spec bs el) l = BlDo a s (Just l) mn tl spec bs el
-  setLabel (BlDoWhile a s _ n tl spec bs el) l = BlDoWhile a s (Just l) n tl spec bs el
+  setLabel (BlDo a s st tl spec bs el) l = BlDo a s (setLabel st l) tl spec bs el
+  setLabel (BlDoWhile a s st tl spec bs el) l = BlDoWhile a s (setLabel st l) tl spec bs el
   setLabel b _ = b
 
 data ProgramUnitName =
@@ -949,6 +968,7 @@ instance Out a => Out (AllocOpt a)
 instance Out UnaryOp
 instance Out BinaryOp
 instance Out a => Out (ForallHeader a)
+instance Out a => Out (BlockConstructStart a)
 
 -- Classifiers on statement and blocks ASTs
 
@@ -1042,3 +1062,4 @@ instance NFData BinaryOp
 instance NFData Only
 instance NFData ModuleNature
 instance NFData Intent
+instance NFData a => NFData (BlockConstructStart a)
diff --git a/src/Language/Fortran/Analysis.hs b/src/Language/Fortran/Analysis.hs
@@ -335,10 +335,10 @@ computeAllLhsVars = concatMap lhsOfStmt . universeBi
 -- | Set of expressions used -- not defined -- by an AST-block.
 blockRhsExprs :: Data a => Block a -> [Expression a]
 blockRhsExprs (BlStatement _ _ _ s) = statementRhsExprs s
-blockRhsExprs (BlDo _ _ _ _ _ (Just (DoSpecification _ _ (StExpressionAssign _ _ lhs rhs) e1 e2)) _ _)
+blockRhsExprs (BlDo _ _ _ _ (Just (DoSpecification _ _ (StExpressionAssign _ _ lhs rhs) e1 e2)) _ _)
   | ExpSubscript _ _ _ subs <- lhs = universeBi (rhs, e1, e2) ++ universeBi subs
   | otherwise                      = universeBi (rhs, e1, e2)
-blockRhsExprs (BlDoWhile _ _ e1 _ _ e2 _ _) = universeBi (e1, e2)
+blockRhsExprs (BlDoWhile _ _ e1 _ e2 _ _) = universeBi (e1, e2)
 blockRhsExprs (BlIf _ _ e1 _ e2 _ _)        = universeBi (e1, e2)
 blockRhsExprs b                             = universeBi b
 
@@ -360,7 +360,7 @@ blockVarUses :: forall a. Data a => Block (Analysis a) -> [Name]
 blockVarUses (BlStatement _ _ _ (StExpressionAssign _ _ lhs rhs))
   | ExpSubscript _ _ _ subs <- lhs = allVars rhs ++ concatMap allVars (aStrip subs)
   | otherwise                      = allVars rhs
-blockVarUses (BlDo _ _ _ _ _ (Just (DoSpecification _ _ (StExpressionAssign _ _ lhs rhs) e1 e2)) _ _)
+blockVarUses (BlDo _ _ _ _ (Just (DoSpecification _ _ (StExpressionAssign _ _ lhs rhs) e1 e2)) _ _)
   | ExpSubscript _ _ _ subs <- lhs = allVars rhs ++ allVars e1 ++ maybe [] allVars e2 ++ concatMap allVars (aStrip subs)
   | otherwise                      = allVars rhs ++ allVars e1 ++ maybe [] allVars e2
 blockVarUses (BlStatement _ _ _ st@StDeclaration{}) = concat [ rhsOfDecls d | d <- universeBi st ]
@@ -372,14 +372,14 @@ blockVarUses (BlStatement _ _ _ st@StDeclaration{}) = concat [ rhsOfDecls d | d
 blockVarUses (BlStatement _ _ _ (StCall _ _ f@(ExpValue _ _ (ValIntrinsic _)) _))
   | Just uses <- intrinsicUses f = uses
 blockVarUses (BlStatement _ _ _ (StCall _ _ _ (Just aexps))) = allVars aexps
-blockVarUses (BlDoWhile _ _ e1 _ _ e2 _ _) = maybe [] allVars e1 ++ allVars e2
+blockVarUses (BlDoWhile _ _ e1 _ e2 _ _) = maybe [] allVars (Just e1) ++ allVars e2
 blockVarUses (BlIf _ _ e1 _ e2 _ _)        = maybe [] allVars e1 ++ concatMap (maybe [] allVars) e2
 blockVarUses b                             = allVars b
 
 -- | Set of names defined by an AST-block.
 blockVarDefs :: Data a => Block (Analysis a) -> [Name]
 blockVarDefs b@BlStatement{} = allLhsVars b
-blockVarDefs (BlDo _ _ _ _ _ (Just doSpec) _ _)  = allLhsVarsDoSpec doSpec
+blockVarDefs (BlDo _ _ _ _ (Just doSpec) _ _)  = allLhsVarsDoSpec doSpec
 blockVarDefs _                      = []
 
 -- form name: n[i]

diff --git a/src/Language/Fortran/Analysis/BBlocks.hs b/src/Language/Fortran/Analysis/BBlocks.hs
@@ -91,11 +91,12 @@ labelWithinBlocks = perBlock'
     perBlock' :: Block (Analysis a) -> Block (Analysis a)
     perBlock' b =
       case b of
+        -- TODO
         BlStatement a s e st               -> BlStatement a s (mfill i e) (fill i st)
         BlIf        a s e1 mn e2 bss el    -> BlIf        a s (mfill i e1) mn (mmfill i e2) bss el
         BlCase      a s e1 mn e2 is bss el -> BlCase      a s (mfill i e1) mn (fill i e2) (mmfill i is) bss el
-        BlDo        a s e1 mn tl e2 bs el  -> BlDo        a s (mfill i e1) mn tl (mfill i e2) bs el
-        BlDoWhile   a s e1 n tl e2 bs el   -> BlDoWhile   a s (mfill i e1) n tl (fill i e2) bs el
+        BlDo        a s e1 tl e2 bs el     -> BlDo        a s (fill i e1) tl (mfill i e2) bs el
+        BlDoWhile   a s e1 tl e2 bs el     -> BlDoWhile   a s (fill i e1) tl (fill i e2) bs el
         _                             -> b
       where i = insLabel $ getAnnotation b
 
@@ -407,14 +408,14 @@ perBlock b@(BlStatement a ss _ (StIfLogical _ _ exp stm)) = do
 perBlock b@(BlStatement _ _ _ StIfArithmetic{}) =
   -- Treat an arithmetic if similarly to a goto
   processLabel b >> addToBBlock b >> closeBBlock_
-perBlock b@(BlDo _ _ _ _ _ (Just spec) bs _) = do
+perBlock b@(BlDo _ _ _ _ (Just spec) bs _) = do
   let DoSpecification _ _ (StExpressionAssign _ _ _ e1) e2 me3 = spec
   _  <- processFunctionCalls e1
   _  <- processFunctionCalls e2
   _  <- case me3 of Just e3 -> Just `fmap` processFunctionCalls e3; Nothing -> return Nothing
   perDoBlock Nothing b bs
-perBlock b@(BlDo _ _ _ _ _ Nothing bs _) = perDoBlock Nothing b bs
-perBlock b@(BlDoWhile _ _ _ _ _ exp bs _) = perDoBlock (Just exp) b bs
+perBlock b@(BlDo _ _ _ _ Nothing bs _) = perDoBlock Nothing b bs
+perBlock b@(BlDoWhile _ _ _ _ exp bs _) = perDoBlock (Just exp) b bs
 perBlock b@(BlStatement _ _ _ StReturn{}) =
   processLabel b >> addToBBlock b >> closeBBlock_
 perBlock b@(BlStatement _ _ _ StGotoUnconditional{}) =
@@ -555,8 +556,8 @@ genTemp str = do
 -- Strip nested code not necessary since it is duplicated in another
 -- basic block.
 stripNestedBlocks :: Block a -> Block a
-stripNestedBlocks (BlDo a s l mn tl ds _ el)     = BlDo a s l mn tl ds [] el
-stripNestedBlocks (BlDoWhile a s l tl n e _ el)  = BlDoWhile a s l tl n e [] el
+stripNestedBlocks (BlDo a s x tl ds _ el)        = BlDo a s x tl ds [] el
+stripNestedBlocks (BlDoWhile a s x n e _ el)  = BlDoWhile a s x n e [] el
 stripNestedBlocks (BlIf a s l mn exps _ el)      = BlIf a s l mn exps [] el
 stripNestedBlocks (BlCase a s l mn sc inds _ el) = BlCase a s l mn sc inds [] el
 stripNestedBlocks b                              = b
@@ -784,13 +785,13 @@ showBlock (BlStatement _ _ mlab st)
         StExit{}                     -> "exit"
         _                            -> "<unhandled statement: " ++ show (toConstr (fmap (const ()) st)) ++ ">"
 showBlock (BlIf _ _ mlab _ (Just e1:_) _ _) = showLab mlab ++ "if " ++ showExpr e1 ++ "\\l"
-showBlock (BlDo _ _ mlab _ _ (Just spec) _ _) =
-    showLab mlab ++ "do " ++ showExpr e1 ++ " <- " ++
+showBlock (BlDo _ _ x _ (Just spec) _ _) =
+    showLab (getLabel x) ++ "do " ++ showExpr e1 ++ " <- " ++
       showExpr e2 ++ ", " ++
       showExpr e3 ++ ", " ++
       maybe "1" showExpr me4 ++ "\\l"
   where DoSpecification _ _ (StExpressionAssign _ _ e1 e2) e3 me4 = spec
-showBlock (BlDo _ _ _ _ _ Nothing _ _) = "do"
+showBlock (BlDo _ _ _ _ Nothing _ _) = "do"
 showBlock (BlComment{})                = ""
 showBlock b = "<unhandled block: " ++ show (toConstr (fmap (const ()) b)) ++ ">"
 

diff --git a/src/Language/Fortran/PrettyPrint.hs b/src/Language/Fortran/PrettyPrint.hs
@@ -286,7 +286,7 @@ instance IndentablePretty (Block a) where
         abstract | v >= Fortran2003 && abstractp = "abstract "
                  | otherwise = empty
 
-    pprint v (BlDo _ _ mLabel mn tl doSpec body el) i
+    pprint v (BlDo _ _ (BlockConstructStart _ _ mLabel mn) tl doSpec body el) i
       | v >= Fortran77Extended =
         labeledIndent mLabel
           (pprint' v mn <?> colon <+>
@@ -310,7 +310,7 @@ instance IndentablePretty (Block a) where
             then indent i (pprint' v label <+> stDoc)
             else pprint' v mLabel `overlay` indent i stDoc
 
-    pprint v (BlDoWhile _ _ mLabel mName mTarget cond body el) i
+    pprint v (BlDoWhile _ _ (BlockConstructStart _ _ mLabel mName) mTarget cond body el) i
        | v >= Fortran77Extended =
         labeledIndent mLabel
           (pprint' v mName <?> colon <+>

diff --git a/src/Language/Fortran/Transformation/Grouping.hs b/src/Language/Fortran/Transformation/Grouping.hs
@@ -85,17 +85,19 @@ groupDo' (b:bs) = b' : bs'
     (b', bs') = case b of
       BlStatement a s label st
         -- Do While statement
-        | StDoWhile _ _ mTarget Nothing condition <- st ->
+        | StDoWhile a' s' mName Nothing condition <- st ->
           let ( blocks, leftOverBlocks, endLabel, stEnd ) =
-                collectNonDoBlocks groupedBlocks mTarget
-          in ( BlDoWhile a (getTransSpan s stEnd) label mTarget Nothing condition blocks endLabel
+                collectNonDoBlocks groupedBlocks mName
+              blStartStmt = BlockConstructStart a' s' label mName
+           in ( BlDoWhile a (getTransSpan s stEnd) blStartStmt Nothing condition blocks endLabel
              , leftOverBlocks)
         -- Vanilla do statement
-        | StDo _ _ mName Nothing doSpec <- st ->
+        | StDo a' s' mName Nothing doSpec <- st ->
           let ( blocks, leftOverBlocks, endLabel, stEnd ) =
                 collectNonDoBlocks groupedBlocks mName
-          in ( BlDo a (getTransSpan s stEnd) label mName Nothing doSpec blocks endLabel
-             , leftOverBlocks)
+              blStartStmt = BlockConstructStart a' s' label mName
+           in ( BlDo a (getTransSpan s stEnd) blStartStmt Nothing doSpec blocks endLabel
+              , leftOverBlocks)
       b'' | containsGroups b'' ->
         ( applyGroupingToSubblocks groupDo' b'', groupedBlocks )
       _ -> ( b, groupedBlocks )
@@ -137,16 +139,18 @@ groupLabeledDo' (b:bs) = b' : bs'
   where
     (b', bs') = case b of
       BlStatement a s label
-        (StDo _ _ mn tl@Just{} doSpec) ->
+        (StDo a' ss' mn tl@Just{} doSpec) ->
           let ( blocks, leftOverBlocks, lastLabel ) =
                 collectNonLabeledDoBlocks tl groupedBlocks
-          in ( BlDo a (getTransSpan s blocks) label mn tl doSpec blocks lastLabel
+              blStartStmt = BlockConstructStart a' ss' label mn
+          in ( BlDo a (getTransSpan s blocks) blStartStmt tl doSpec blocks lastLabel
              , leftOverBlocks )
       BlStatement a s label
-        (StDoWhile _ _ mn tl@Just{} cond) ->
+        (StDoWhile a' ss' mn tl@Just{} cond) ->
           let ( blocks, leftOverBlocks, lastLabel ) =
                 collectNonLabeledDoBlocks tl groupedBlocks
-          in ( BlDoWhile a (getTransSpan s blocks) label mn tl cond blocks lastLabel
+              blStartStmt = BlockConstructStart a' ss' label mn
+          in ( BlDoWhile a (getTransSpan s blocks) blStartStmt tl cond blocks lastLabel
              , leftOverBlocks )
       b'' | containsGroups b'' ->
         ( applyGroupingToSubblocks groupLabeledDo' b'', groupedBlocks )
@@ -210,10 +214,10 @@ applyGroupingToSubblocks f b
     BlIf a s l mn conds (map f blocks) el
   | BlCase a s l mn scrutinee conds blocks         el <- b =
     BlCase a s l mn scrutinee conds (map f blocks) el
-  | BlDo a s l n tl doSpec blocks     el <- b =
-    BlDo a s l n tl doSpec (f blocks) el
-  | BlDoWhile a s l n tl doSpec blocks     el <- b =
-    BlDoWhile a s l n tl doSpec (f blocks) el
+  | BlDo a s x tl doSpec blocks     el <- b =
+    BlDo a s x tl doSpec (f blocks) el
+  | BlDoWhile a s x tl doSpec blocks     el <- b =
+    BlDoWhile a s x tl doSpec (f blocks) el
   | BlInterface{} <- b =
       error "Interface blocks do not have groupable subblocks. Must not occur."
   | BlComment{} <- b =

diff --git a/test/Language/Fortran/PrettyPrintSpec.hs b/test/Language/Fortran/PrettyPrintSpec.hs
@@ -303,7 +303,7 @@ spec =
       describe "Do While" $ do
         it "prints simple do while loop" $ do
           let cond = ExpBinary () u LFA.GT (varGen "i") (intGen 42)
-          let bl = BlDoWhile () u Nothing (Just "my_block") Nothing cond body Nothing
+          let bl = BlDoWhile () u (BlockConstructStart () u Nothing (Just "my_block")) Nothing cond body Nothing
           let expect = unlines [ "my_block: do while ((i > 42))"
                                , "print *, i"
                                , "i = (i - 1)"
@@ -312,7 +312,7 @@ spec =
 
         it "prints a labelled do while loop" $ do
           let cond = ExpBinary () u LFA.GT (varGen "i") (intGen 42)
-          let bl = BlDoWhile () u Nothing Nothing (Just (intGen 10)) cond body Nothing
+          let bl = BlDoWhile () u (BlockConstructStart () u Nothing Nothing) (Just (intGen 10)) cond body Nothing
           let expect = unlines [ "do 10 while ((i > 42))"
                                , "print *, i"
                                , "i = (i - 1)" ]
@@ -323,23 +323,23 @@ spec =
         let doSpec = DoSpecification () u iAssign (intGen 9) (Just (intGen 2))
 
         it "prints 90 style do loop" $ do
-          let bl = BlDo () u Nothing Nothing Nothing (Just doSpec) body Nothing
+          let bl = BlDo () u (BlockConstructStart () u Nothing Nothing) Nothing (Just doSpec) body Nothing
           let expect = unlines [ "do i = 1, 9, 2"
                                , "print *, i"
                                , "i = (i - 1)"
                                , "end do" ]
           pprint Fortran90 bl Nothing `shouldBe` text expect
 
         it "prints named infinite do loop" $ do
-          let bl = BlDo () u Nothing (Just "joker") Nothing Nothing body Nothing
+          let bl = BlDo () u (BlockConstructStart () u Nothing (Just "joker")) Nothing Nothing body Nothing
           let expect = unlines [ "joker: do"
                                , "print *, i"
                                , "i = (i - 1)"
                                , "end do joker" ]
           pprint Fortran90 bl Nothing `shouldBe` text expect
 
         it "prints named labeled do loop" $ do
-          let bl = BlDo () u Nothing (Just "joker") (Just $ intGen 42) (Just doSpec) body (Just $ intGen 42)
+          let bl = BlDo () u (BlockConstructStart () u Nothing (Just "joker")) (Just $ intGen 42) (Just doSpec) body (Just $ intGen 42)
           let expect = unlines [ "joker: do 42 i = 1, 9, 2"
                                , "print *, i"
                                , "i = (i - 1)"
@@ -348,7 +348,7 @@ spec =
 
         it "prints vanilla labeled do loop" $ do
           let body2 = body ++ [ BlStatement () u (Just $ intGen 42) (StContinue () u) ]
-          let bl = BlDo () u Nothing Nothing (Just $ intGen 42) (Just doSpec) body2 (Just $ intGen 42)
+          let bl = BlDo () u (BlockConstructStart () u Nothing Nothing) (Just $ intGen 42) (Just doSpec) body2 (Just $ intGen 42)
           let expect = unlines [ "      do 42 i = 1, 9, 2"
                                , "        print *, i"
                                , "        i = (i - 1)"