diff --git a/lib/inference.ml b/lib/inference.ml
index 2ddb9b6..54e9702 100644
--- a/lib/inference.ml
+++ b/lib/inference.ml
@@ -1,11 +1,7 @@
 (** Infer the type of a given term and, if exists, returns the type of the term *)
 let typeof t =
   let rec infer env = function
-    | Term.Var id ->
-      ( (match TypeSubstitution.find id env with
-         | Some ty -> Some ty
-         | None -> Some (Type.Var id))
-      , env )
+    | Term.Var id -> TypeSubstitution.find id env, env
     | Term.IntConst _ -> Some Type.Int, env
     | Term.Binop (t1, _, t2) ->
       (* Both operands must have type Int *)
@@ -28,7 +24,8 @@ let typeof t =
           | Term.Second -> Some ty2, env)
        | _ -> None, env)
     | Term.Fun (id, body) ->
-      (match infer env body with
+      let arg = TypeSubstitution.singleton id (Type.Var id) in
+      (match infer (TypeSubstitution.compose env arg) body with
        | Some ty_body, env' ->
          (match infer env' (Term.Var id) with
           | Some ty, _ -> Some (Type.Arrow (ty, ty_body)), env'