Infer range types

Author: oxalica

crates/ra_hir_def/src/body/lower.rs

@@ -8,7 +8,7 @@ use hir_expand::{
 use ra_arena::Arena;
 use ra_syntax::{
     ast::{
-        self, ArgListOwner, ArrayExprKind, LiteralKind, LoopBodyOwner, NameOwner,
+        self, ArgListOwner, ArrayExprKind, LiteralKind, LoopBodyOwner, NameOwner, RangeOp,
         TypeAscriptionOwner,
     },
     AstNode, AstPtr,
@@ -429,10 +429,28 @@ where
                 let index = self.collect_expr_opt(e.index());
                 self.alloc_expr(Expr::Index { base, index }, syntax_ptr)
             }
+            ast::Expr::RangeExpr(e) => {
+                let lhs = e.start().map(|lhs| self.collect_expr(lhs));
+                let rhs = e.end().map(|rhs| self.collect_expr(rhs));
+                match (lhs, e.op_kind(), rhs) {
+                    (None, _, None) => self.alloc_expr(Expr::RangeFull, syntax_ptr),
+                    (Some(lhs), _, None) => self.alloc_expr(Expr::RangeFrom { lhs }, syntax_ptr),
+                    (None, Some(RangeOp::Inclusive), Some(rhs)) => {
+                        self.alloc_expr(Expr::RangeToInclusive { rhs }, syntax_ptr)
+                    }
+                    (Some(lhs), Some(RangeOp::Inclusive), Some(rhs)) => {
+                        self.alloc_expr(Expr::RangeInclusive { lhs, rhs }, syntax_ptr)
+                    }
+                    // If RangeOp is missing, fallback to exclusive range.
+                    (None, _, Some(rhs)) => self.alloc_expr(Expr::RangeTo { rhs }, syntax_ptr),
+                    (Some(lhs), _, Some(rhs)) => {
+                        self.alloc_expr(Expr::Range { lhs, rhs }, syntax_ptr)
+                    }
+                }
+            }
 
             // FIXME implement HIR for these:
             ast::Expr::Label(_e) => self.alloc_expr(Expr::Missing, syntax_ptr),
-            ast::Expr::RangeExpr(_e) => self.alloc_expr(Expr::Missing, syntax_ptr),
             ast::Expr::MacroCall(e) => match self.expander.enter_expand(self.db, e) {
                 Some((mark, expansion)) => {
                     let id = self.collect_expr(expansion);

crates/ra_hir_def/src/expr.rs

@@ -130,6 +130,24 @@ pub enum Expr {
         rhs: ExprId,
         op: Option<BinaryOp>,
     },
+    RangeFull,
+    RangeFrom {
+        lhs: ExprId,
+    },
+    RangeTo {
+        rhs: ExprId,
+    },
+    Range {
+        lhs: ExprId,
+        rhs: ExprId,
+    },
+    RangeToInclusive {
+        rhs: ExprId,
+    },
+    RangeInclusive {
+        lhs: ExprId,
+        rhs: ExprId,
+    },
     Index {
         base: ExprId,
         index: ExprId,
@@ -284,15 +302,21 @@ impl Expr {
             Expr::Lambda { body, .. } => {
                 f(*body);
             }
-            Expr::BinaryOp { lhs, rhs, .. } => {
+            Expr::BinaryOp { lhs, rhs, .. }
+            | Expr::Range { lhs, rhs }
+            | Expr::RangeInclusive { lhs, rhs } => {
                 f(*lhs);
                 f(*rhs);
             }
             Expr::Index { base, index } => {
                 f(*base);
                 f(*index);
             }
-            Expr::Field { expr, .. }
+            Expr::RangeFull => {}
+            Expr::RangeFrom { lhs: expr }
+            | Expr::RangeTo { rhs: expr }
+            | Expr::RangeToInclusive { rhs: expr }
+            | Expr::Field { expr, .. }
             | Expr::Await { expr }
             | Expr::Try { expr }
             | Expr::Cast { expr, .. }

crates/ra_hir_def/src/path.rs

@@ -409,6 +409,36 @@ pub mod known {
         Path::from_simple_segments(PathKind::Abs, vec![name::STD, name::OPS, name::TRY_TYPE])
     }
 
+    pub fn std_ops_range() -> Path {
+        Path::from_simple_segments(PathKind::Abs, vec![name::STD, name::OPS, name::RANGE_TYPE])
+    }
+
+    pub fn std_ops_range_from() -> Path {
+        Path::from_simple_segments(PathKind::Abs, vec![name::STD, name::OPS, name::RANGE_FROM_TYPE])
+    }
+
+    pub fn std_ops_range_full() -> Path {
+        Path::from_simple_segments(PathKind::Abs, vec![name::STD, name::OPS, name::RANGE_FULL_TYPE])
+    }
+
+    pub fn std_ops_range_inclusive() -> Path {
+        Path::from_simple_segments(
+            PathKind::Abs,
+            vec![name::STD, name::OPS, name::RANGE_INCLUSIVE_TYPE],
+        )
+    }
+
+    pub fn std_ops_range_to() -> Path {
+        Path::from_simple_segments(PathKind::Abs, vec![name::STD, name::OPS, name::RANGE_TO_TYPE])
+    }
+
+    pub fn std_ops_range_to_inclusive() -> Path {
+        Path::from_simple_segments(
+            PathKind::Abs,
+            vec![name::STD, name::OPS, name::RANGE_TO_INCLUSIVE_TYPE],
+        )
+    }
+
     pub fn std_result_result() -> Path {
         Path::from_simple_segments(PathKind::Abs, vec![name::STD, name::RESULT, name::RESULT_TYPE])
     }

crates/ra_hir_expand/src/name.rs

@@ -140,6 +140,12 @@ pub const RESULT_TYPE: Name = Name::new_inline_ascii(6, b"Result");
 pub const OUTPUT_TYPE: Name = Name::new_inline_ascii(6, b"Output");
 pub const TARGET_TYPE: Name = Name::new_inline_ascii(6, b"Target");
 pub const BOX_TYPE: Name = Name::new_inline_ascii(3, b"Box");
+pub const RANGE_FROM_TYPE: Name = Name::new_inline_ascii(9, b"RangeFrom");
+pub const RANGE_FULL_TYPE: Name = Name::new_inline_ascii(9, b"RangeFull");
+pub const RANGE_INCLUSIVE_TYPE: Name = Name::new_inline_ascii(14, b"RangeInclusive");
+pub const RANGE_TO_INCLUSIVE_TYPE: Name = Name::new_inline_ascii(16, b"RangeToInclusive");
+pub const RANGE_TO_TYPE: Name = Name::new_inline_ascii(7, b"RangeTo");
+pub const RANGE_TYPE: Name = Name::new_inline_ascii(5, b"Range");
 
 // Builtin Macros
 pub const FILE_MACRO: Name = Name::new_inline_ascii(4, b"file");

crates/ra_hir_ty/src/infer.rs

@@ -577,6 +577,42 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
         let struct_ = self.resolver.resolve_known_struct(self.db, &path)?;
         Some(struct_.into())
     }
+
+    fn resolve_range_full(&self) -> Option<AdtId> {
+        let path = known::std_ops_range_full();
+        let struct_ = self.resolver.resolve_known_struct(self.db, &path)?;
+        Some(struct_.into())
+    }
+
+    fn resolve_range(&self) -> Option<AdtId> {
+        let path = known::std_ops_range();
+        let struct_ = self.resolver.resolve_known_struct(self.db, &path)?;
+        Some(struct_.into())
+    }
+
+    fn resolve_range_inclusive(&self) -> Option<AdtId> {
+        let path = known::std_ops_range_inclusive();
+        let struct_ = self.resolver.resolve_known_struct(self.db, &path)?;
+        Some(struct_.into())
+    }
+
+    fn resolve_range_from(&self) -> Option<AdtId> {
+        let path = known::std_ops_range_from();
+        let struct_ = self.resolver.resolve_known_struct(self.db, &path)?;
+        Some(struct_.into())
+    }
+
+    fn resolve_range_to(&self) -> Option<AdtId> {
+        let path = known::std_ops_range_to();
+        let struct_ = self.resolver.resolve_known_struct(self.db, &path)?;
+        Some(struct_.into())
+    }
+
+    fn resolve_range_to_inclusive(&self) -> Option<AdtId> {
+        let path = known::std_ops_range_to_inclusive();
+        let struct_ = self.resolver.resolve_known_struct(self.db, &path)?;
+        Some(struct_.into())
+    }
 }
 
 /// The ID of a type variable.

crates/ra_hir_ty/src/infer/expr.rs

@@ -415,6 +415,47 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                 }
                 _ => Ty::Unknown,
             },
+            Expr::RangeFull => match self.resolve_range_full() {
+                Some(adt) => Ty::simple(TypeCtor::Adt(adt)),
+                None => Ty::Unknown,
+            },
+            Expr::Range { lhs, rhs } => {
+                let lhs_ty = self.infer_expr(*lhs, &Expectation::none());
+                let rhs_ty = self.infer_expr(*rhs, &Expectation::has_type(lhs_ty));
+                match self.resolve_range() {
+                    Some(adt) => Ty::apply_one(TypeCtor::Adt(adt), rhs_ty),
+                    None => Ty::Unknown,
+                }
+            }
+            Expr::RangeInclusive { lhs, rhs } => {
+                let lhs_ty = self.infer_expr(*lhs, &Expectation::none());
+                let rhs_ty = self.infer_expr(*rhs, &Expectation::has_type(lhs_ty));
+                match self.resolve_range_inclusive() {
+                    Some(adt) => Ty::apply_one(TypeCtor::Adt(adt), rhs_ty),
+                    None => Ty::Unknown,
+                }
+            }
+            Expr::RangeFrom { lhs } => {
+                let ty = self.infer_expr(*lhs, &Expectation::none());
+                match self.resolve_range_from() {
+                    Some(adt) => Ty::apply_one(TypeCtor::Adt(adt), ty),
+                    None => Ty::Unknown,
+                }
+            }
+            Expr::RangeTo { rhs } => {
+                let ty = self.infer_expr(*rhs, &Expectation::none());
+                match self.resolve_range_to() {
+                    Some(adt) => Ty::apply_one(TypeCtor::Adt(adt), ty),
+                    None => Ty::Unknown,
+                }
+            }
+            Expr::RangeToInclusive { rhs } => {
+                let ty = self.infer_expr(*rhs, &Expectation::none());
+                match self.resolve_range_to_inclusive() {
+                    Some(adt) => Ty::apply_one(TypeCtor::Adt(adt), ty),
+                    None => Ty::Unknown,
+                }
+            }
             Expr::Index { base, index } => {
                 let _base_ty = self.infer_expr(*base, &Expectation::none());
                 let _index_ty = self.infer_expr(*index, &Expectation::none());

crates/ra_hir_ty/src/tests.rs

@@ -221,6 +221,56 @@ mod collections {
     assert_eq!("&str", type_at_pos(&db, pos));
 }
 
+#[test]
+fn infer_ranges() {
+    let (db, pos) = TestDB::with_position(
+        r#"
+//- /main.rs crate:main deps:std
+fn test() {
+    let a = ..;
+    let b = 1..;
+    let c = ..2u32;
+    let d = 1..2usize;
+    let e = ..=10;
+    let f = 'a'..='z';
+
+    let t = (a, b, c, d, e, f);
+    t<|>;
+}
+
+//- /std.rs crate:std
+#[prelude_import] use prelude::*;
+mod prelude {}
+
+pub mod ops {
+    pub struct Range<Idx> {
+        pub start: Idx,
+        pub end: Idx,
+    }
+    pub struct RangeFrom<Idx> {
+        pub start: Idx,
+    }
+    struct RangeFull;
+    pub struct RangeInclusive<Idx> {
+        start: Idx,
+        end: Idx,
+        is_empty: u8,
+    }
+    pub struct RangeTo<Idx> {
+        pub end: Idx,
+    }
+    pub struct RangeToInclusive<Idx> {
+        pub end: Idx,
+    }
+}
+"#,
+    );
+    assert_eq!(
+        "(RangeFull, RangeFrom<i32>, RangeTo<u32>, Range<usize>, RangeToInclusive<i32>, RangeInclusive<char>)",
+        type_at_pos(&db, pos),
+    );
+}
+
 #[test]
 fn infer_while_let() {
     let (db, pos) = TestDB::with_position(