From f370f05e63a29c2c0c65efe90ba948714a0d6cdd Mon Sep 17 00:00:00 2001
From: Vidush Singhal <vidush.sl@gmail.com>
Date: Tue, 3 Dec 2024 13:28:01 -0500
Subject: [PATCH] generate soa regions in inferlocations

---
 .../src/Gibbon/Passes/InferLocations.hs       | 34 +++++++++++++++----
 1 file changed, 27 insertions(+), 7 deletions(-)

diff --git a/gibbon-compiler/src/Gibbon/Passes/InferLocations.hs b/gibbon-compiler/src/Gibbon/Passes/InferLocations.hs
index ee4fc4c4..306b2684 100644
--- a/gibbon-compiler/src/Gibbon/Passes/InferLocations.hs
+++ b/gibbon-compiler/src/Gibbon/Passes/InferLocations.hs
@@ -142,8 +142,8 @@ convertFunTy ddefs (from,to,isPar) = do
     from' <- mapM (convertTy ddefs useSoA) from
     to'   <- convertTy ddefs useSoA to
     -- For this simple version, we assume every location is in a separate region:
-    lrm1 <- concat <$> mapM (toLRM Input) from'
-    lrm2 <- toLRM Output to'
+    lrm1 <- concat <$> mapM (toLRM useSoA Input) from'
+    lrm2 <- toLRM useSoA Output to'
     dbgTraceIt "convertFunTy: " dbgTraceIt (sdoc (from', to', lrm1, lrm2, useSoA)) dbgTraceIt "\n" return $ ArrowTy2 { locVars = lrm1 ++ lrm2
                      , arrIns  = from'
                      , arrEffs = S.empty
@@ -151,10 +151,31 @@ convertFunTy ddefs (from,to,isPar) = do
                      , locRets = []
                      , hasParallelism = isPar }
  where
-   toLRM md ls =
-       mapM (\v -> do r <- freshLocVar "r"
-                      return $ LRM v (AoSR $ VarR (unwrapLocVar r)) md)
-            (F.toList ls)
+   toLRM soa md ls  = do
+                      case soa of 
+                          True -> mapM (\v -> do
+                                              dataBufRegion <- freshLocVar "r"
+                                              case v of 
+                                                Single _ -> error "InferLocations : toLRM : Expected an SoA location!\n"
+                                                SoA _ fieldLocs -> do
+                                                                   fieldRegions <- getSoARegionsFromLocs fieldLocs 
+                                                                   let region = SoAR (VarR (unwrapLocVar dataBufRegion)) fieldRegions
+                                                                   return $ LRM v region md
+                                       ) (F.toList ls)
+                          False -> mapM (\v -> do r <- freshLocVar "r"
+                                                  return $ LRM v (AoSR $ VarR (unwrapLocVar r)) md) (F.toList ls)
+
+getSoARegionsFromLocs :: [((DataCon, Int), Var)] -> PassM [((DataCon, Int), Region)]
+getSoARegionsFromLocs locs = case locs of 
+                                  [] -> return [] 
+                                  (a, b):rst -> do 
+                                                regionVariable <- freshLocVar "r"
+                                                let region = VarR (unwrapLocVar regionVariable)
+                                                rst' <- getSoARegionsFromLocs rst
+                                                let elem = (a, region) 
+                                                return $ [elem] ++ rst'
+
+
 
 convertTy :: DDefs1 -> Bool -> Ty1 -> PassM Ty2
 convertTy ddefs useSoA ty = case useSoA of 
@@ -162,7 +183,6 @@ convertTy ddefs useSoA ty = case useSoA of
                             True ->  case ty of 
                                         PackedTy tycon _ -> do 
                                                              dconBuff <- freshLocVar "loc"
-                                                             let ddef = lookupDDef ddefs tycon
                                                              let dcons = getConOrdering ddefs tycon
                                                              locsForFields <- convertTyHelperSoAParent tycon ddefs dcons
                                                              let soaLocation = SoA (unwrapLocVar dconBuff) locsForFields