Rain: added a pass for finding and tagging the main function

RainPassTest.hs

@@ -228,6 +228,39 @@ testRecordInfNames3 = testPassShouldFail "testRecordInfNames3" (recordInfNameTyp
   where
     orig = A.Rep m (A.ForEach m (simpleName "c") (intLiteral 0)) skipP
 
+--Easy way to string two passes together; creates a pass-like function that applies the left-hand pass then the right-hand pass.  Associative.
+(>>>) :: Monad m => (a -> m b) -> (b -> m c) -> a -> m c
+(>>>) f0 f1 x = (f0 x) >>= f1
+
+testFindMain0 :: Test
+testFindMain0 = testPassWithStateCheck "testFindMain0" orig ((uniquifyAndResolveVars >>> findMain) orig) (return ()) check
+  where
+    orig = A.Spec m (A.Specification m (A.Name m A.ProcName "main") $ A.Proc m A.PlainSpec [] (A.Skip m)) $ A.Several m []
+    check state = assertEqual "testFindMain0" [("main",(A.Name m A.ProcName "main"))] (csMainLocals state)
+
+testFindMain1 :: Test
+testFindMain1 = testPassWithStateCheck "testFindMain1" orig ((uniquifyAndResolveVars >>> findMain) orig) (return ()) check
+  where
+    orig = A.Spec m (A.Specification m (A.Name m A.ProcName "foo") $ A.Proc m A.PlainSpec [] (A.Skip m)) $ A.Several m []
+    check state = assertEqual "testFindMain1" [] (csMainLocals state)
+    
+testFindMain2 :: Test
+testFindMain2 = testPassWithStateCheck "testFindMain2" orig ((uniquifyAndResolveVars >>> findMain) orig) (return ()) check
+  where
+    orig = A.Spec m (A.Specification m (A.Name m A.ProcName "main") $ A.Proc m A.PlainSpec [] (A.Skip m)) $ 
+             A.Spec m (A.Specification m (A.Name m A.ProcName "foo") $ A.Proc m A.PlainSpec [] (A.Skip m)) $
+               A.Several m []
+    check state = assertEqual "testFindMain2" [("main",(A.Name m A.ProcName "main"))] (csMainLocals state)
+
+testFindMain3 :: Test
+testFindMain3 = testPassWithStateCheck "testFindMain3" orig ((uniquifyAndResolveVars >>> findMain) orig) (return ()) check
+  where
+    orig = A.Spec m (A.Specification m (A.Name m A.ProcName "foo") $ A.Proc m A.PlainSpec [] (A.Skip m)) $ 
+             A.Spec m (A.Specification m (A.Name m A.ProcName "main") $ A.Proc m A.PlainSpec [] (A.Skip m)) $
+               A.Several m []
+    check state = assertEqual "testFindMain3" [("main",(A.Name m A.ProcName "main"))] (csMainLocals state)
+
+
 --Returns the list of tests:
 tests :: Test
 tests = TestList
@@ -242,6 +275,10 @@ tests = TestList
    ,testRecordInfNames1
    ,testRecordInfNames2
    ,testRecordInfNames3
+   ,testFindMain0
+   ,testFindMain1
+   ,testFindMain2
+   ,testFindMain3
  ]
 
 

RainPasses.hs

@@ -22,6 +22,8 @@ import TestUtil
 import qualified AST as A
 import Pass
 import Data.Generics
+import qualified Data.Map as Map
+import Control.Monad.State
 import Types
 import CompState
 import Errors
@@ -30,7 +32,8 @@ rainPasses :: [(String,Pass)]
 rainPasses = 
      [ ("Resolve Int -> Int64",transformInt)
        ,("Uniquify variable declarations, record declared types and resolve variable names",uniquifyAndResolveVars)
-       ,("Record inferred name types in dictionary",recordInfNameTypes)
+       ,("Record inferred name types in dictionary",recordInfNameTypes) --depends on uniquifyAndResolveVars
+       ,("Find and tag the main function",findMain) --depends on uniquifyAndResolveVars
        ,("Convert seqeach/pareach loops into classic replicated SEQ/PAR",transformEach)
      ]
 
@@ -93,6 +96,18 @@ recordInfNameTypes = everywhereM (mkM recordInfNameTypes')
            return input
     recordInfNameTypes' r = return r
 
+findMain :: Data t => t -> PassM t
+--Because findMain runs after uniquifyAndResolveVars, the types of all the process will have been recorded
+--Therefore this pass doesn't actually need to walk the tree, it just has to look for a process named "main"
+--in the CompState, and pull it out into csMainLocals
+findMain x = do st <- get
+                put (findMain' st)
+                return x
+  where
+    findMain' :: CompState -> CompState 
+    findMain' st = case (Map.lookup "main" (csNames st)) of
+      Just n -> st {csMainLocals = [("main",A.Name {A.nameName = "main", A.nameMeta = A.ndMeta n, A.nameType = A.ndNameType n})]}
+      Nothing -> st 
 
 transformEach :: Data t => t -> PassM t
 transformEach = everywhereM (mkM transformEach')