From 8651be33055e31517525595577c4b440f9e4ec39 Mon Sep 17 00:00:00 2001
From: Markus Westerlind <marwes91@gmail.com>
Date: Sat, 17 Mar 2018 01:32:59 +0100
Subject: [PATCH] fix: Take implicits into account in type queries on infix
 expressions

---
 base/src/ast.rs       | 12 ++------
 base/src/types/mod.rs | 18 ++++++++++++
 base/tests/types.rs   | 67 ++++++++++++++++++++++++++++++++++++++++++-
 repl/tests/rexpect.rs |  2 +-
 4 files changed, 87 insertions(+), 12 deletions(-)

diff --git a/base/src/ast.rs b/base/src/ast.rs
index abda84e7b8..7bd10011be 100644
--- a/base/src/ast.rs
+++ b/base/src/ast.rs
@@ -819,15 +819,7 @@ impl Typed for Expr<Symbol> {
             | Expr::Tuple { ref typ, .. } => Ok(typ.clone()),
             Expr::Literal(ref lit) => lit.try_type_of(env),
             Expr::IfElse(_, ref arm, _) => arm.try_type_of(env),
-            Expr::Infix { ref op, .. } => op.value
-                .typ
-                .as_function()
-                .and_then(|(_, ret)| ret.as_function())
-                .map(|(_, ret)| ret.clone())
-                .ok_or_else(|| {
-                    debug!("`{}` is not a binary function type", op.value.typ);
-                    "Expected function type in binop".to_string()
-                }),
+            Expr::Infix { ref op, .. } => get_return_type(env, &op.value.typ, 2),
             Expr::LetBindings(_, ref expr)
             | Expr::TypeBindings(_, ref expr)
             | Expr::Do(Do { body: ref expr, .. }) => expr.try_type_of(env),
@@ -878,7 +870,7 @@ fn get_return_type(
         return Ok(alias_type.clone());
     }
 
-    let alias_type = alias_type.remove_forall();
+    let alias_type = alias_type.remove_forall_and_implicit_args();
     if let Some((_, ret)) = alias_type.as_function() {
         return get_return_type(env, ret, arg_count - 1);
     }
diff --git a/base/src/types/mod.rs b/base/src/types/mod.rs
index b74e5bb329..51f1871f74 100644
--- a/base/src/types/mod.rs
+++ b/base/src/types/mod.rs
@@ -765,6 +765,16 @@ where
         })
     }
 
+    pub fn function_implicit<I>(args: I, ret: T) -> T
+    where
+        I: IntoIterator<Item = T>,
+        I::IntoIter: DoubleEndedIterator<Item = T>,
+    {
+        args.into_iter().rev().fold(ret, |body, arg| {
+            T::from(Type::Function(ArgType::Implicit, arg, body))
+        })
+    }
+
     pub fn generic(typ: Generic<Id>) -> T {
         T::from(Type::Generic(typ))
     }
@@ -1094,6 +1104,14 @@ impl<Id> ArcType<Id> {
         remove_forall(self)
     }
 
+    pub fn remove_forall_and_implicit_args(&self) -> &ArcType<Id> {
+        match **self {
+            Type::Function(ArgType::Implicit, _, ref typ) => typ.remove_forall_and_implicit_args(),
+            Type::Forall(_, ref typ, _) => typ.remove_forall_and_implicit_args(),
+            _ => self,
+        }
+    }
+
     pub fn skolemize(&self, named_variables: &mut FnvMap<Id, ArcType<Id>>) -> ArcType<Id>
     where
         Id: Clone + Eq + Hash,
diff --git a/base/tests/types.rs b/base/tests/types.rs
index c1d0ea7398..bbbd1c23c7 100644
--- a/base/tests/types.rs
+++ b/base/tests/types.rs
@@ -7,9 +7,12 @@ use std::ops::Deref;
 
 use pretty::{Arena, DocAllocator};
 
-use base::kind::Kind;
+use base::ast::{Expr, Literal, SpannedExpr, Typed, TypedIdent};
+use base::kind::{ArcKind, Kind, KindEnv};
 use base::types::*;
+use base::symbol::{Symbol, SymbolRef};
 use base::source::Source;
+use base::pos::{self, BytePos, Span, Spanned};
 
 fn type_con<I, T>(s: I, args: Vec<T>) -> Type<I, T>
 where
@@ -341,3 +344,65 @@ fn break_record() {
 "#
     );
 }
+
+pub struct MockEnv;
+
+impl KindEnv for MockEnv {
+    fn find_kind(&self, _id: &SymbolRef) -> Option<ArcKind> {
+        None
+    }
+}
+
+impl TypeEnv for MockEnv {
+    fn find_type(&self, _id: &SymbolRef) -> Option<&ArcType> {
+        None
+    }
+
+    fn find_type_info(&self, _id: &SymbolRef) -> Option<&Alias<Symbol, ArcType>> {
+        None
+    }
+
+    fn find_record(
+        &self,
+        _fields: &[Symbol],
+        _selector: RecordSelector,
+    ) -> Option<(ArcType, ArcType)> {
+        None
+    }
+}
+
+pub type SpExpr = SpannedExpr<Symbol>;
+
+pub fn intern(s: &str) -> Symbol {
+    Symbol::from(s)
+}
+
+pub fn no_loc<T>(value: T) -> Spanned<T, BytePos> {
+    pos::spanned(Span::default(), value)
+}
+
+pub fn int(i: i64) -> SpExpr {
+    no_loc(Expr::Literal(Literal::Int(i)))
+}
+
+pub fn binop(l: SpExpr, s: &str, r: SpExpr) -> SpExpr {
+    no_loc(Expr::Infix {
+        lhs: Box::new(l),
+        op: no_loc(TypedIdent::new(intern(s))),
+        rhs: Box::new(r),
+        implicit_args: Vec::new(),
+    })
+}
+
+#[test]
+fn take_implicits_into_account_on_infix_type() {
+    let mut expr = binop(int(1), "+", int(2));
+    if let Expr::Infix { ref mut op, .. } = expr.value {
+        op.value.typ = Type::function_implicit(
+            vec![Type::int()],
+            Type::function(vec![Type::int(), Type::int()], Type::int()),
+        );
+    }
+
+    assert_eq!(expr.env_type_of(&MockEnv), Type::int());
+}
diff --git a/repl/tests/rexpect.rs b/repl/tests/rexpect.rs
index 4847019052..2fe9c12133 100644
--- a/repl/tests/rexpect.rs
+++ b/repl/tests/rexpect.rs
@@ -92,7 +92,7 @@ fn expression_types() {
     let mut repl = REPL::new();
 
     repl.test(":t 5", Some("Int"));
-    // repl.test(":t 5 + 5", Some("Int -> Int"));
+    repl.test(":t 5 + 5", Some("Int"));
     repl.test(":t \"gluon\"", Some("String"));
 }