feat(session): refactor states according to new interface

2 files changed, 18 insertions, 16 deletions

diff --git a/lib/session.ml b/lib/session.ml
index 4db13ff..7e1052f 100644
--- a/lib/session.ml
+++ b/lib/session.ml
@@ -4,15 +4,16 @@ open Lwt_react
 open Stream
 
 type step =
-  | Starting_stream
-  | Negotiating_feature of Feature.requirement * features
+  | Feature of Feature.requirement * features
   | Logged_in of features
 
 type state =
   | Disconnected
   | Connecting of Portal.domain
-  (* TCP/WebSocket connected, not connected in XMPP-land *)
-  | Connected of Portal.t * step
+  (* TCP/WebSocket connected. *)
+  | Connected of Portal.socket
+  (* Stream negotiation *)
+  | Negotiating of Portal.t * step
 
 (** An XMPP session. This type contains a signal representing the state of an XMPP
     connection, and its update function. *)
@@ -34,29 +35,30 @@ let create (config : config) : t Lwt.t =
     | Feature.Mechanisms _ | STARTTLS -> true
     | _ -> false
   and features_next_state = function
-    | feature :: rest -> Negotiating_feature (feature, rest)
+    | feature :: rest -> Feature (feature, rest)
     (* No features returned by stream start. Connection is completed. *)
     | [] -> Logged_in []
   and eq s1 s2 =  (* TODO: move this closer to the state type *)
     match s1, s2 with
-    | Connected (_, step1), Connected (_, step2) -> step1 = step2
+    | Negotiating (_, step1), Negotiating (_, step2) -> step1 = step2
     | _ -> s1 = s2
   in
   let state, update = S.create ~eq Disconnected in
   let connection_map = function
     | Connecting domain ->
-       let+ portal = Portal.connect domain
-       in update (Connected (portal, Starting_stream))
-    | Connected (portal, Starting_stream) ->
+       let+ socket = Portal.connect domain
+       in update (Connected socket)
+    | Connected socket ->
+       let* portal = Portal.stream socket in
        let+ features = Stream.start portal
        in let next_state = features_next_state features
-          in update (Connected (portal, next_state))
-    | Connected (portal, Negotiating_feature (feature, features)) ->
+          in update (Negotiating (portal, next_state))
+    | Negotiating (portal, Feature (feature, features)) ->
        let+ () = negotiate feature portal config
        in let next_state = if needs_restart (Feature.unwrap feature)
-                           then Starting_stream
-                           else features_next_state features
-          in update (Connected (portal, next_state))
+                           then Connected portal._socket
+                           else Negotiating (portal, features_next_state features)
+          in update next_state
     | _ -> Lwt.return_unit
   in let+ () = S.map_s connection_map state >|= S.keep;
      in { state; update }
diff --git a/test/hello.ml b/test/hello.ml
index 3504af4..9388b75 100644
--- a/test/hello.ml
+++ b/test/hello.ml
@@ -18,8 +18,8 @@ let main =
   let waiter, wakener = Lwt.wait () in
   try%lwt let* { state; update } = create config in
              S.map (function
-                 | Connected (portal, Logged_in _) -> portal.push None;
-                                                      Lwt.wakeup wakener ()
+                 | Negotiating (portal, Logged_in _) -> portal.push None;
+                                                        Lwt.wakeup wakener ()
                  | _ -> ())
                state |> S.keep;
              connect { state; update } config.sasl.jid.domainpart;