Change format of internal data structure for group storage

ianhinder · Nov 21, 2012 · 7a654b1 · 7a654b1
1 parent 2d7502d
commit 7a654b1
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 13 deletions.
diff --git a/Tools/CodeGen/Schedule.m b/Tools/CodeGen/Schedule.m
@@ -28,11 +28,10 @@
    Scheduling
    -------------------------------------------------------------------------- *)
 
-storageStructure[groupName_, timelevels_, param_String] := 
+storageStructure[groupName_, timelevels_] := 
 {
   Group -> groupName, 
-  Timelevels -> timelevels,
-  TimelevelsParameter -> param
+  Timelevels -> timelevels
 };
 
 groupsSetInCalc[calc_, groups_] :=
@@ -274,15 +273,15 @@
             {tl},
             tl = NonevolvedTimelevels[groupFromName[#, groups]];
             If[tl===1,
-               storageStructure[#, tl, "other_timelevels"],
-               storageStructure[#, evolutionTimelevels, "timelevels"]]] &,
+               storageStructure[#, {"other_timelevels", tl}],
+               storageStructure[#, {"timelevels", evolutionTimelevels}]]] &,
           (* over *)
           nonevolvedGroups],
 
-        Map[storageStructure[#, evolutionTimelevels, "timelevels"] &,
+        Map[storageStructure[#, {"timelevels", evolutionTimelevels}] &,
             evolvedGroups],
 
-        Map[storageStructure[#, evolutionTimelevels, "rhs_timelevels"] &,
+        Map[storageStructure[#, {"rhs_timelevels", evolutionTimelevels}] &,
             rhsGroups]];
 
     (* Schedule groups defined in calculations *)

diff --git a/Tools/CodeGen/Thorn.m b/Tools/CodeGen/Thorn.m
@@ -364,12 +364,26 @@
    structure for inclusion in the schedule.ccl file to allocate
    storage for this group. *)
 groupStorage[spec_] :=
-  If[mapContains[spec, TimelevelsParameter],
-     Flatten[Table[{"if (", lookup[spec, TimelevelsParameter], " == ", i, ")\n",
-                    "{\n",
-                    "  STORAGE: ", lookup[spec, Group], "[", i, "]\n",
-                    "}\n"}, {i, 1, lookup[spec, Timelevels]}], 1],
-     {"STORAGE: ", lookup[spec, Group], "[", lookup[spec, Timelevels], "]\n"}]
+  Module[
+    {tls = lookup[spec,Timelevels],
+     group = lookup[spec, Group]},
+    Which[
+      IntegerQ[tls],
+      {"STORAGE: ", group, "[", tls, "]\n"},
+
+      ListQ[tls],
+      If[!MatchQ[tls, {_String, _Integer}],
+         Error["Unrecognized Timelevels value "<>ToString[tls]]];
+      Module[
+        {param,max},
+        {param,max} = tls;
+        Flatten[
+          Table[{"if (", param, " == ", i, ")\n",
+                 "{\n",
+                 "  STORAGE: ", group, "[", i, "]\n",
+                 "}\n"}, {i, 1, max}], 1]],
+
+      True, Error["Unrecognized Timelevels value "<>ToString[tls]]]];
 
 
 (* Given a function scheduling specification as defined above, return