1 files changed, 19 insertions, 19 deletions

diff --git a/tree.sml b/tree.sml
index 0a788af..c97edfb 100644
--- a/tree.sml
+++ b/tree.sml
@@ -82,28 +82,28 @@ structure Tree: TREE = struct
     assemble' n buf
   end
 
-  fun lookup' _ _ Empty _ _ g = (g, NONE)
-    | lookup' buf cmp (Node (k', v', left, right)) k f g =
-        case cmp k k' of
-          LESS => lookup' (Left (k', v', right) :: buf) cmp left k f g
-        | GREATER => lookup' (Right (k', v', left) :: buf) cmp right k f g
-        | EQUAL =>
-        let
-          val (newV, result) = f v'
-        in
-          case newV of
-               NONE => (result, NONE)
-             | SOME v => (result, SOME (assemble (Node (k', v, left, right)) buf))
-        end
-
-  fun lookup2 cmp t k (f, g) =
+  fun lookup' _ _ Empty k f =
   let
-    val (result, newTree) = lookup' [] cmp t k f g
+    val (res, newV) = f NONE
   in
-    (result, case newTree of
-         NONE => t
-       | SOME t => t)
+    case newV of
+      NONE => (res, Empty)
+    | SOME v => (res, Node (k, v, Empty, Empty))
   end
+    | lookup' buf cmp (T as Node (k', v', left, right)) k f =
+      case cmp k k' of
+        LESS => lookup' (Left (k', v', right) :: buf) cmp left k f
+      | GREATER => lookup' (Right (k', v', left) :: buf) cmp right k f
+      | EQUAL =>
+      let
+        val (res, newV) = f $ SOME v'
+      in
+        case newV of
+           NONE => (res, T)
+         | SOME v => (res, assemble (Node (k', v, left, right)) buf)
+      end
+
+  fun lookup2 cmp t k f = lookup' [] cmp t k f
 
   fun print t key2str value2str =
   let