Make Scriptcheck catch more script/message-passing uses, and move the phase earlier in compilation

18 files changed, 119 insertions, 184 deletions

diff --git a/src/cjr.sml b/src/cjr.sml
index c348d01a..3a37b26f 100644
--- a/src/cjr.sml
+++ b/src/cjr.sml
@@ -128,10 +128,7 @@ datatype decl' =
 
 withtype decl = decl' located
 
-datatype sidedness =
-         ServerOnly
-       | ServerAndPull
-       | ServerAndPullAndPush
+datatype sidedness = datatype Mono.sidedness
 
 datatype effect = datatype Export.effect
 datatype export_kind = datatype Export.export_kind
diff --git a/src/cjrize.sml b/src/cjrize.sml
index 9e41fda4..0f4bdb42 100644
--- a/src/cjrize.sml
+++ b/src/cjrize.sml
@@ -694,7 +694,7 @@ fun cifyDecl ((d, loc), sm) =
       | L.DPolicy _ => (NONE, NONE, sm)
       | L.DOnError n => (SOME (L'.DOnError n, loc), NONE, sm)
 
-fun cjrize ds =
+fun cjrize (ds, sideInfo) =
     let
         val (dsF, ds, ps, sm) = foldl (fn (d, (dsF, ds, ps, sm)) =>
                                           let
@@ -722,6 +722,13 @@ fun cjrize ds =
                                               (dsF, ds, ps, Sm.clearDeclares sm)
                                           end)
                                       ([], [], [], Sm.empty) ds
+
+        val sideInfo = foldl (fn ((n, mode), mp) => IM.insert (mp, n, mode)) IM.empty sideInfo
+
+        val ps = map (fn (ek, s, n, ts, t, _, b) =>
+                         (ek, s, n, ts, t,
+                          getOpt (IM.find (sideInfo, n), L'.ServerOnly),
+                          b)) ps
     in
         (List.revAppend (dsF, rev ds),
          ps)
diff --git a/src/compiler.sig b/src/compiler.sig
index 7e4f2f6a..fcf664eb 100644
--- a/src/compiler.sig
+++ b/src/compiler.sig
@@ -116,12 +116,12 @@ signature COMPILER = sig
     val mono_shake : (Mono.file, Mono.file) phase
     val iflow : (Mono.file, Mono.file) phase
     val namejs : (Mono.file, Mono.file) phase
+    val scriptcheck : (Mono.file, Mono.file) phase
     val jscomp : (Mono.file, Mono.file) phase
     val fuse : (Mono.file, Mono.file) phase
     val pathcheck : (Mono.file, Mono.file) phase
     val sidecheck : (Mono.file, Mono.file) phase
     val cjrize : (Mono.file, Cjr.file) phase
-    val scriptcheck : (Cjr.file, Cjr.file) phase
     val prepare : (Cjr.file, Cjr.file) phase
     val checknest : (Cjr.file, Cjr.file) phase
     val sqlify : (Mono.file, Cjr.file) phase
@@ -170,6 +170,7 @@ signature COMPILER = sig
     val toIflow : (string, Mono.file) transform
     val toNamejs : (string, Mono.file) transform
     val toNamejs_untangle : (string, Mono.file) transform
+    val toScriptcheck : (string, Mono.file) transform
     val toJscomp : (string, Mono.file) transform
     val toMono_opt3 : (string, Mono.file) transform
     val toFuse : (string, Mono.file) transform
@@ -184,7 +185,6 @@ signature COMPILER = sig
     val toPathcheck : (string, Mono.file) transform
     val toSidecheck : (string, Mono.file) transform
     val toCjrize : (string, Cjr.file) transform
-    val toScriptcheck : (string, Cjr.file) transform
     val toPrepare : (string, Cjr.file) transform
     val toChecknest : (string, Cjr.file) transform
     val toSqlify : (string, Cjr.file) transform
diff --git a/src/compiler.sml b/src/compiler.sml
index f8dd07e2..77542811 100644
--- a/src/compiler.sml
+++ b/src/compiler.sml
@@ -1363,12 +1363,19 @@ val toNamejs = transform namejs "namejs" o toIflow
 
 val toNamejs_untangle = transform untangle "namejs_untangle" o toNamejs
 
+val scriptcheck = {
+    func = ScriptCheck.classify,
+    print = MonoPrint.p_file MonoEnv.empty
+}
+
+val toScriptcheck = transform scriptcheck "scriptcheck" o toNamejs_untangle
+
 val jscomp = {
     func = JsComp.process,
     print = MonoPrint.p_file MonoEnv.empty
 }
 
-val toJscomp = transform jscomp "jscomp" o toNamejs_untangle
+val toJscomp = transform jscomp "jscomp" o toScriptcheck
 
 val toMono_opt3 = transform mono_opt "mono_opt3" o toJscomp
 
@@ -1410,19 +1417,12 @@ val cjrize = {
 
 val toCjrize = transform cjrize "cjrize" o toSidecheck
 
-val scriptcheck = {
-    func = ScriptCheck.classify,
-    print = CjrPrint.p_file CjrEnv.empty
-}
-
-val toScriptcheck = transform scriptcheck "scriptcheck" o toCjrize
-
 val prepare = {
     func = Prepare.prepare,
     print = CjrPrint.p_file CjrEnv.empty
 }
 
-val toPrepare = transform prepare "prepare" o toScriptcheck
+val toPrepare = transform prepare "prepare" o toCjrize
 
 val checknest = {
     func = fn f => if #supportsNestedPrepared (Settings.currentDbms ()) then f else Checknest.annotate f,
diff --git a/src/fuse.sml b/src/fuse.sml
index 565fc591..5193e59a 100644
--- a/src/fuse.sml
+++ b/src/fuse.sml
@@ -144,9 +144,9 @@ fun fuse file =
                  (funcs, maxName))
             end
 
-        val (file, _) = ListUtil.foldlMap doDecl (IM.empty, U.File.maxName file + 1) file
+        val (ds, _) = ListUtil.foldlMap doDecl (IM.empty, U.File.maxName file + 1) (#1 file)
     in
-        file
+        (ds, #2 file)
     end
 
 end
diff --git a/src/iflow.sml b/src/iflow.sml
index fe0be731..8c933dc4 100644
--- a/src/iflow.sml
+++ b/src/iflow.sml
@@ -1795,7 +1795,7 @@ fun evalExp env (e as (_, loc)) k =
 
 datatype var_source = Input of int | SubInput of int | Unknown
 
-fun check file =
+fun check (file : file) =
     let
         val () = (St.reset ();
                   rfuns := IM.empty)
@@ -1810,7 +1810,7 @@ fun check file =
         val exptd = foldl (fn ((d, _), exptd) =>
                               case d of
                                   DExport (_, _, n, _, _, _) => IS.add (exptd, n)
-                                | _ => exptd) IS.empty file
+                                | _ => exptd) IS.empty (#1 file)
 
         fun decl (d, loc) =
             case d of
@@ -2071,7 +2071,7 @@ fun check file =
                                         
               | _ => ()
     in
-        app decl file
+        app decl (#1 file)
     end
 
 val check = fn file =>
diff --git a/src/jscomp.sml b/src/jscomp.sml
index ea34a3b5..ffb68ab2 100644
--- a/src/jscomp.sml
+++ b/src/jscomp.sml
@@ -61,7 +61,7 @@ exception CantEmbed of typ
 
 fun inString {needle, haystack} = String.isSubstring needle haystack
 
-fun process file =
+fun process (file : file) =
     let
         val (someTs, nameds) =
             foldl (fn ((DVal (_, n, t, e, _), _), (someTs, nameds)) => (someTs, IM.insert (nameds, n, e))
@@ -77,7 +77,7 @@ fun process file =
                                      someTs) someTs dts,
                        nameds)
                     | (_, state) => state)
-                  (IM.empty, IM.empty) file
+                  (IM.empty, IM.empty) (#1 file)
 
         fun str loc s = (EPrim (Prim.String s), loc)
 
@@ -1304,7 +1304,7 @@ fun process file =
                         listInjectors = TM.empty,
                         decoders = IM.empty,
                         maxName = U.File.maxName file + 1}
-                       file
+                       (#1 file)
 
         val inf = TextIO.openIn (OS.Path.joinDirFile {dir = Settings.libJs (), file = "urweb.js"})
         fun lines acc =
@@ -1334,7 +1334,7 @@ fun process file =
                 ""
     in
         TextIO.closeIn inf;
-        (DJavaScript script, ErrorMsg.dummySpan) :: ds
+        ((DJavaScript script, ErrorMsg.dummySpan) :: ds, #2 file)
     end
 
 end
diff --git a/src/mono.sml b/src/mono.sml
index 4a0278fd..f269c52d 100644
--- a/src/mono.sml
+++ b/src/mono.sml
@@ -157,6 +157,11 @@ datatype decl' =
 
 withtype decl = decl' located
 
-type file = decl list
+datatype sidedness =
+         ServerOnly
+       | ServerAndPull
+       | ServerAndPullAndPush
+
+type file = decl list * (int * sidedness) list
 
 end
diff --git a/src/mono_print.sml b/src/mono_print.sml
index e5ef4cf8..12b36f2a 100644
--- a/src/mono_print.sml
+++ b/src/mono_print.sml
@@ -530,7 +530,7 @@ fun p_decl env (dAll as (d, _) : decl) =
                           p_policy env p]
       | DOnError _ => string "ONERROR"
                           
-fun p_file env file =
+fun p_file env (file, _) =
     let
         val (pds, _) = ListUtil.foldlMap (fn (d, env) =>
                                              (p_decl env d,
diff --git a/src/mono_reduce.sml b/src/mono_reduce.sml
index 71c87095..e7fac5ed 100644
--- a/src/mono_reduce.sml
+++ b/src/mono_reduce.sml
@@ -308,7 +308,7 @@ val freeInAbs = U.Exp.existsB {typ = fn _ => false,
                                              U.Exp.RelE _ => n + 1
                                            | _ => n} 0
 
-fun reduce file =
+fun reduce (file : file) =
     let
         val (timpures, impures, absCounts) =
             foldl (fn ((d, _), (timpures, impures, absCounts)) =>
@@ -366,7 +366,7 @@ fun reduce file =
                                      absCounts vis)
                             | _ => (timpures, impures, absCounts)
                       end)
-                  (IS.empty, IS.empty, IM.empty) file
+                  (IS.empty, IS.empty, IM.empty) (#1 file)
 
         val uses = U.File.fold {typ = fn (_, m) => m,
                                 exp = fn (e, m) =>
diff --git a/src/mono_shake.sml b/src/mono_shake.sml
index b6de9410..5818fea0 100644
--- a/src/mono_shake.sml
+++ b/src/mono_shake.sml
@@ -41,7 +41,7 @@ type free = {
      exp : IS.set
 }
 
-fun shake file =
+fun shake (file : file) =
     let
         val (cdef, edef) = foldl (fn ((DDatatype dts, _), (cdef, edef)) =>
                                      (foldl (fn ((_, n, xncs), cdef) => IM.insert (cdef, n, xncs)) cdef dts, edef)
@@ -60,7 +60,7 @@ fun shake file =
                                    | ((DTask _, _), acc) => acc
                                    | ((DPolicy _, _), acc) => acc
                                    | ((DOnError _, _), acc) => acc)
-                                 (IM.empty, IM.empty) file
+                                 (IM.empty, IM.empty) (#1 file)
 
         fun typ (c, s) =
             case c of
@@ -130,7 +130,7 @@ fun shake file =
                         usedVars st e1
                     end
                   | ((DOnError n, _), (page_cs, page_es)) => (page_cs, IS.add (page_es, n))
-                  | (_, st) => st) (IS.empty, IS.empty) file
+                  | (_, st) => st) (IS.empty, IS.empty) (#1 file)
 
         val s = {con = page_cs, exp = page_es}
 
@@ -145,20 +145,20 @@ fun shake file =
                                  NONE => raise Fail "MonoShake: Couldn't find 'val'"
                                | SOME (t, e) => shakeExp s e) s page_es
     in
-        List.filter (fn (DDatatype dts, _) => List.exists (fn (_, n, _) => IS.member (#con s, n)) dts
-                      | (DVal (_, n, _, _, _), _) => IS.member (#exp s, n)
-                      | (DValRec vis, _) => List.exists (fn (_, n, _, _, _) => IS.member (#exp s, n)) vis
-                      | (DExport _, _) => true
-                      | (DTable _, _) => true
-                      | (DSequence _, _) => true
-                      | (DView _, _) => true
-                      | (DDatabase _, _) => true
-                      | (DJavaScript _, _) => true
-                      | (DCookie _, _) => true
-                      | (DStyle _, _) => true
-                      | (DTask _, _) => true
-                      | (DPolicy _, _) => true
-                      | (DOnError _, _) => true) file
+        (List.filter (fn (DDatatype dts, _) => List.exists (fn (_, n, _) => IS.member (#con s, n)) dts
+                       | (DVal (_, n, _, _, _), _) => IS.member (#exp s, n)
+                       | (DValRec vis, _) => List.exists (fn (_, n, _, _, _) => IS.member (#exp s, n)) vis
+                       | (DExport _, _) => true
+                       | (DTable _, _) => true
+                       | (DSequence _, _) => true
+                       | (DView _, _) => true
+                       | (DDatabase _, _) => true
+                       | (DJavaScript _, _) => true
+                       | (DCookie _, _) => true
+                       | (DStyle _, _) => true
+                       | (DTask _, _) => true
+                       | (DPolicy _, _) => true
+                       | (DOnError _, _) => true) (#1 file), #2 file)
     end
 
 end
diff --git a/src/mono_util.sml b/src/mono_util.sml
index 58498996..61638858 100644
--- a/src/mono_util.sml
+++ b/src/mono_util.sml
@@ -664,9 +664,9 @@ fun mapfoldB (all as {bind, ...}) =
     let
         val mfd = Decl.mapfoldB all
 
-        fun mff ctx ds =
+        fun mff ctx (ds, ps) =
             case ds of
-                nil => S.return2 nil
+                nil => S.return2 (nil, ps)
               | d :: ds' =>
                 S.bind2 (mfd ctx d,
                          fn d' =>
@@ -705,9 +705,9 @@ fun mapfoldB (all as {bind, ...}) =
                                       | DPolicy _ => ctx
                                       | DOnError _ => ctx
                             in
-                                S.map2 (mff ctx' ds',
-                                     fn ds' =>
-                                        d' :: ds')
+                                S.map2 (mff ctx' (ds', ps),
+                                     fn (ds', _) =>
+                                        (d' :: ds', ps))
                             end)
     in
         mff
@@ -741,27 +741,28 @@ fun fold {typ, exp, decl} s d =
         S.Continue (_, s) => s
       | S.Return _ => raise Fail "MonoUtil.File.fold: Impossible"
 
-val maxName = foldl (fn ((d, _) : decl, count) =>
-                        case d of
-                            DDatatype dts =>
-                            foldl (fn ((_, n, ns), count) =>
-                                      foldl (fn ((_, n', _), m) => Int.max (n', m))
-                                            (Int.max (n, count)) ns) count dts
-                          | DVal (_, n, _, _, _) => Int.max (n, count)
-                          | DValRec vis => foldl (fn ((_, n, _, _, _), count) => Int.max (n, count)) count vis
-                          | DExport _ => count
-                          | DTable _ => count
-                          | DSequence _ => count
-                          | DView _ => count
-                          | DDatabase _ => count
-                          | DJavaScript _ => count
-                          | DCookie _ => count
-                          | DStyle _ => count
-                          | DTask _ => count
-                          | DPolicy _ => count
-                          | DOnError _ => count) 0
-
-fun appLoc f =
+fun maxName (f : file) =
+    foldl (fn ((d, _) : decl, count) =>
+              case d of
+                  DDatatype dts =>
+                  foldl (fn ((_, n, ns), count) =>
+                            foldl (fn ((_, n', _), m) => Int.max (n', m))
+                                  (Int.max (n, count)) ns) count dts
+                | DVal (_, n, _, _, _) => Int.max (n, count)
+                | DValRec vis => foldl (fn ((_, n, _, _, _), count) => Int.max (n, count)) count vis
+                | DExport _ => count
+                | DTable _ => count
+                | DSequence _ => count
+                | DView _ => count
+                | DDatabase _ => count
+                | DJavaScript _ => count
+                | DCookie _ => count
+                | DStyle _ => count
+                | DTask _ => count
+                | DPolicy _ => count
+                | DOnError _ => count) 0 (#1 f)
+
+fun appLoc f (fl : file) =
     let
         val eal = Exp.appLoc f
 
@@ -790,7 +791,7 @@ fun appLoc f =
               | PolUpdate e1 => eal e1
               | PolSequence e1 => eal e1
     in
-        app appl
+        app appl (#1 fl)
     end
 
 end
diff --git a/src/monoize.sml b/src/monoize.sml
index e07c0c90..ce7bfbe9 100644
--- a/src/monoize.sml
+++ b/src/monoize.sml
@@ -4656,7 +4656,7 @@ fun monoize env file =
         pvars := RM.empty;
         pvarDefs := [];
         pvarOldDefs := [];
-        rev ds
+        (rev ds, [])
     end
 
 end
diff --git a/src/name_js.sml b/src/name_js.sml
index 70ac000c..53abd7a3 100644
--- a/src/name_js.sml
+++ b/src/name_js.sml
@@ -72,7 +72,7 @@ fun squish vs = U.Exp.mapB {typ = fn x => x,
 
 fun rewrite file =
     let
-        val (file, _) = ListUtil.foldlMapConcat (fn (d, nextName) =>
+        val (ds, _) = ListUtil.foldlMapConcat (fn (d, nextName) =>
                                                     let
                                                         val (d, (nextName, newDs)) =
                                                             U.Decl.foldMapB {typ = fn x => x,
@@ -143,9 +143,9 @@ fun rewrite file =
                                                                       DValRec vis => [(DValRec (vis @ newDs), #2 d)]
                                                                     | _ => List.revAppend (map (fn vi => (DVal vi, #2 d)) newDs, [d]),
                                                          nextName)
-                                                    end) (U.File.maxName file + 1) file
+                                                    end) (U.File.maxName file + 1) (#1 file)
     in
-        file
+        (ds, #2 file)
     end
 
 end
diff --git a/src/pathcheck.sml b/src/pathcheck.sml
index 15405db7..c1bb667b 100644
--- a/src/pathcheck.sml
+++ b/src/pathcheck.sml
@@ -110,6 +110,6 @@ fun checkDecl ((d, loc), (funcs, rels, cookies, styles)) =
           | _ => (funcs, rels, cookies, styles)
     end
 
-fun check ds = ignore (foldl checkDecl (SS.empty, SS.empty, SS.empty, SS.empty) ds)
+fun check (ds, _) = ignore (foldl checkDecl (SS.empty, SS.empty, SS.empty, SS.empty) ds)
 
 end
diff --git a/src/scriptcheck.sig b/src/scriptcheck.sig
index bc9b6377..afb557b7 100644
--- a/src/scriptcheck.sig
+++ b/src/scriptcheck.sig
@@ -27,6 +27,6 @@
 
 signature SCRIPT_CHECK = sig
 
-    val classify : Cjr.file -> Cjr.file
+    val classify : Mono.file -> Mono.file
 
 end
diff --git a/src/scriptcheck.sml b/src/scriptcheck.sml
index 6c6c5588..e5db476a 100644
--- a/src/scriptcheck.sml
+++ b/src/scriptcheck.sml
@@ -27,7 +27,7 @@
 
 structure ScriptCheck :> SCRIPT_CHECK = struct
 
-open Cjr
+open Mono
 
 structure SS = BinarySetFn(struct
                            type ord_key = string
@@ -35,98 +35,31 @@ structure SS = BinarySetFn(struct
                            end)
 structure IS = IntBinarySet
 
-val pullBasis = SS.addList (SS.empty,
-                            ["new_client_source",
-                             "get_client_source",
-                             "set_client_source"])
-
 val pushBasis = SS.addList (SS.empty,
                             ["new_channel",
                              "self"])
 
-val events = ["abort",
-              "blur",
-              "change",
-              "click",
-              "dblclick",
-              "error",
-              "focus",
-              "keydown",
-              "keypress",
-              "keyup",
-              "load",
-              "mousedown",
-              "mousemove",
-              "mouseout",
-              "mouseover",
-              "mouseup",
-              "reset",
-              "resize",
-              "select",
-              "submit",
-              "unload"]
-                
-val scriptWords = "<script"
-                   :: map (fn s => " on" ^ s ^ "='") events
-
-val pushWords = ["rv("]
-
 fun classify (ds, ps) =
     let
         val proto = Settings.currentProtocol ()
 
         fun inString {needle, haystack} = String.isSubstring needle haystack
 
-        fun hasClient {basis, words, onload} csids =
-            let
-                fun hasClient e =
-                    case #1 e of
-                        EPrim (Prim.String s) => List.exists (fn n => inString {needle = n, haystack = s}) words
-                      | EPrim _ => false
-                      | ERel _ => false
-                      | ENamed n => IS.member (csids, n)
-                      | ECon (_, _, NONE) => false
-                      | ECon (_, _, SOME e) => hasClient e
-                      | ENone _ => false
-                      | ESome (_, e) => hasClient e
-                      | EFfi ("Basis", x) => SS.member (basis, x)
-                      | EFfi _ => false
-                      | EFfiApp ("Basis", "maybe_onload",
-                                 [((EFfiApp ("Basis", "strcat", all as [_, ((EPrim (Prim.String s), _), _)]), _), _)]) =>
-                        List.exists (hasClient o #1) all
-                        orelse (onload andalso size s > 0)
-                      | EFfiApp ("Basis", x, es) => SS.member (basis, x)
-                                                    orelse List.exists (hasClient o #1) es
-                      | EFfiApp (_, _, es) => List.exists (hasClient o #1) es
-                      | EApp (e, es) => hasClient e orelse List.exists hasClient es
-                      | EUnop (_, e) => hasClient e
-                      | EBinop (_, e1, e2) => hasClient e1 orelse hasClient e2
-                      | ERecord (_, xes) => List.exists (hasClient o #2) xes
-                      | EField (e, _) => hasClient e
-                      | ECase (e, pes, _) => hasClient e orelse List.exists (hasClient o #2) pes
-                      | EError (e, _) => hasClient e
-                      | EReturnBlob {blob = e1, mimeType = e2, ...} => hasClient e1 orelse hasClient e2
-                      | ERedirect (e, _) => hasClient e
-                      | EWrite e => hasClient e
-                      | ESeq (e1, e2) => hasClient e1 orelse hasClient e2
-                      | ELet (_, _, e1, e2) => hasClient e1 orelse hasClient e2
-                      | EQuery {query, body, initial, ...} => hasClient query orelse hasClient body
-                                                              orelse hasClient initial
-                      | EDml {dml, ...} => hasClient dml
-                      | ENextval {seq, ...} => hasClient seq
-                      | ESetval {seq, count, ...} => hasClient seq orelse hasClient count
-                      | EUnurlify (e, _, _) => hasClient e
-            in
-                hasClient
-            end
+        fun hasClient {basis, funcs, push} =
+            MonoUtil.Exp.exists {typ = fn _ => false,
+                                 exp = fn ERecv _ => push
+                                        | EFfiApp ("Basis", x, _) => SS.member (basis, x) 
+                                        | EJavaScript _ => not push
+                                        | ENamed n => IS.member (funcs, n)
+                                        | _ => false}
 
         fun decl ((d, _), (pull_ids, push_ids)) =
             let
-                val hasClientPull = hasClient {basis = pullBasis, words = scriptWords, onload = true} pull_ids
-                val hasClientPush = hasClient {basis = pushBasis, words = pushWords, onload = false} push_ids
+                val hasClientPull = hasClient {basis = SS.empty, funcs = pull_ids, push = false}
+                val hasClientPush = hasClient {basis = pushBasis, funcs = push_ids, push = true}
             in
                 case d of
-                    DVal (_, n, _, e) => (if hasClientPull e then
+                    DVal (_, n, _, e, _) => (if hasClientPull e then
                                              IS.add (pull_ids, n)
                                           else
                                               pull_ids,
@@ -134,20 +67,12 @@ fun classify (ds, ps) =
                                               IS.add (push_ids, n)
                                           else
                                               push_ids)
-                  | DFun (_, n, _, _, e) => (if hasClientPull e then
-                                                 IS.add (pull_ids, n)
-                                             else
-                                                 pull_ids,
-                                             if hasClientPush e then
-                                                 IS.add (push_ids, n)
-                                             else
-                                                 push_ids)
-                  | DFunRec xes => (if List.exists (fn (_, _, _, _, e) => hasClientPull e) xes then
+                  | DValRec xes => (if List.exists (fn (_, _, _, e, _) => hasClientPull e) xes then
                                        foldl (fn ((_, n, _, _, _), pull_ids) => IS.add (pull_ids, n))
                                              pull_ids xes
                                     else
                                         pull_ids,
-                                    if List.exists (fn (_, _, _, _, e) => hasClientPush e) xes then
+                                    if List.exists (fn (_, _, _, e, _) => hasClientPush e) xes then
                                         foldl (fn ((_, n, _, _, _), push_ids) => IS.add (push_ids, n))
                                               push_ids xes
                                     else
@@ -159,21 +84,21 @@ fun classify (ds, ps) =
 
         val foundBad = ref false
 
-        val ps = map (fn (ek, x, n, ts, t, _, b) =>
-                         (ek, x, n, ts, t,
-                          if IS.member (push_ids, n) then
-                              (if not (#persistent proto) andalso not (!foundBad) then
-                                   (foundBad := true;
-                                    ErrorMsg.error ("This program needs server push, but the current protocol ("
-                                                    ^ #name proto ^ ") doesn't support that."))
-                               else
-                                   ();
-                               ServerAndPullAndPush)
-                          else if IS.member (pull_ids, n) then
-                              ServerAndPull
-                          else
-                              ServerOnly,
-                          b)) ps
+        val all_ids = IS.union (pull_ids, push_ids)
+
+        val ps = map (fn n =>
+                         (n, if IS.member (push_ids, n) then
+                                 (if not (#persistent proto) andalso not (!foundBad) then
+                                      (foundBad := true;
+                                       ErrorMsg.error ("This program needs server push, but the current protocol ("
+                                                       ^ #name proto ^ ") doesn't support that."))
+                                  else
+                                      ();
+                                  ServerAndPullAndPush)
+                             else if IS.member (pull_ids, n) then
+                                 ServerAndPull
+                             else
+                                 ServerOnly)) (IS.listItems all_ids)
     in
         (ds, ps)
     end
diff --git a/src/untangle.sml b/src/untangle.sml
index 373cfe18..bcb90ed6 100644
--- a/src/untangle.sml
+++ b/src/untangle.sml
@@ -43,7 +43,7 @@ fun exp (e, s) =
 
       | _ => s
 
-fun untangle file =
+fun untangle (file : file) =
     let
         fun decl (dAll as (d, loc)) =
             case d of
@@ -208,7 +208,7 @@ fun untangle file =
                 end
               | _ => [dAll]
     in
-        ListUtil.mapConcat decl file
+        (ListUtil.mapConcat decl (#1 file), #2 file)
     end
 
 end