Nested demo

2 files changed, 40 insertions, 9 deletions

diff --git a/src/core_untangle.sml b/src/core_untangle.sml
index 6f424614..ededded0 100644
--- a/src/core_untangle.sml
+++ b/src/core_untangle.sml
@@ -45,6 +45,15 @@ fun exp (e, s) =
 
 fun untangle file =
     let
+        val edefs = foldl (fn ((d, _), edefs) =>
+                              case d of
+                                  DVal (_, n, _, e, _) => IM.insert (edefs, n, e)
+                                | DValRec vis =>
+                                  foldl (fn ((_, n, _, e, _), edefs) =>
+                                            IM.insert (edefs, n, e)) edefs vis
+                                | _ => edefs)
+                    IM.empty file
+
         fun decl (dAll as (d, loc)) =
             case d of
                 DValRec vis =>
@@ -52,16 +61,35 @@ fun untangle file =
                     val thisGroup = foldl (fn ((_, n, _, _, _), thisGroup) =>
                                               IS.add (thisGroup, n)) IS.empty vis
 
+                    val expUsed = U.Exp.fold {con = default,
+                                              kind = default,
+                                              exp = exp} IS.empty
+
                     val used = foldl (fn ((_, n, _, e, _), used) =>
                                        let
-                                           val usedHere = U.Exp.fold {con = default,
-                                                                      kind = default,
-                                                                      exp = exp} IS.empty e
+                                           val usedHere = expUsed e
                                        in
-                                           IM.insert (used, n, IS.intersection (usedHere, thisGroup))
+                                           IM.insert (used, n, usedHere)
                                        end)
                                IM.empty vis
 
+                    fun expand used =
+                        IS.foldl (fn (n, used) =>
+                                     case IM.find (edefs, n) of
+                                         NONE => used
+                                       | SOME e =>
+                                         let
+                                             val usedHere = expUsed e
+                                         in
+                                             if IS.isEmpty (IS.difference (usedHere, used)) then
+                                                 used
+                                             else
+                                                 expand (IS.union (usedHere, used))
+                                         end)
+                        used used
+
+                    val used = IM.map (fn s => IS.intersection (expand s, thisGroup)) used
+
                     fun p_graph reachable =
                         IM.appi (fn (n, reachableHere) =>
                                     (print (Int.toString n);
diff --git a/src/unnest.sml b/src/unnest.sml
index f226a678..b56daf8a 100644
--- a/src/unnest.sml
+++ b/src/unnest.sml
@@ -137,7 +137,7 @@ fun squishExp (nr, cfv, efv) =
 
 type state = {
      maxName : int,
-     decls : decl list
+     decls : (string * int * con * exp) list
 }
 
 fun kind (k, st) = (k, st)
@@ -278,11 +278,9 @@ fun exp ((ks, ts), e as old, st : state) =
                                               end)
                                           vis
 
-                            val d = (DValRec vis, #2 ed)
-
                             val ts = map (fn (x, _, t, _) => (x, t)) vis @ ts
                         in
-                            ([], (ts, maxName, d :: ds, subs))
+                            ([], (ts, maxName, vis @ ds, subs))
                         end)
                 (ts, #maxName st, #decls st, []) eds
         in
@@ -319,8 +317,13 @@ fun unnest file =
                 fun explore () =
                     let
                         val (d, st) = unnestDecl st all
+
+                        val ds =
+                            case #1 d of
+                                DValRec vis => [(DValRec (vis @ #decls st), #2 d)]
+                              | _ => [(DValRec (#decls st), #2 d), d]
                     in
-                        (rev (d :: #decls st),
+                        (ds,
                          {maxName = #maxName st,
                           decls = []})
                     end