Add IndexRef and IndexMut

src/librustc/middle/lang_items.rs

@@ -230,6 +230,8 @@ lets_do_this! {
     ShlTraitLangItem,                "shl",                     shl_trait;
     ShrTraitLangItem,                "shr",                     shr_trait;
     IndexTraitLangItem,              "index",                   index_trait;
+    IndexMutTraitLangItem,           "index_mut",               index_mut_trait;
+    IndexRefTraitLangItem,           "index_ref",               index_ref_trait;
 
     EqTraitLangItem,                 "eq",                      eq_trait;
     OrdTraitLangItem,                "ord",                     ord_trait;

src/librustc/middle/resolve.rs

@@ -5286,6 +5286,12 @@ impl Resolver {
             ExprIndex(..) => {
                 let i = self.lang_items.index_trait();
                 self.add_fixed_trait_for_expr(expr.id, i);
+
+                let j = self.lang_items.index_mut_trait();
+                self.add_fixed_trait_for_expr(expr.id, j);
+
+                let k = self.lang_items.index_ref_trait();
+                self.add_fixed_trait_for_expr(expr.id, k);
             }
             _ => {
                 // Nothing to do.
@@ -5382,9 +5388,12 @@ impl Resolver {
     fn add_fixed_trait_for_expr(&mut self,
                                     expr_id: NodeId,
                                     trait_id: Option<DefId>) {
+
+        let trait_map = &mut self.trait_map;
+
         match trait_id {
             Some(trait_id) => {
-                self.trait_map.insert(expr_id, @RefCell::new(~[trait_id]));
+                trait_map.mangle(expr_id, trait_id, |_, id| @RefCell::new(~[id]), |_, old, id| old.borrow_mut().get().push(id));
             }
             None => {}
         }

src/librustc/middle/typeck/check/mod.rs

@@ -1969,6 +1969,21 @@ pub fn check_expr_with_unifier(fcx: @FnCtxt,
         fcx.write_ty(id, if_ty);
     }
 
+    fn has_op_method(fcx: @FnCtxt,
+                        callee_id: ast::NodeId,
+                        op_ex: &ast::Expr,
+                        self_t: ty::t,
+                        opname: ast::Name,
+                        args: &[@ast::Expr],
+                        deref_args: DerefArgs,
+                        autoderef_receiver: AutoderefReceiverFlag,
+                        _expected_result: Option<ty::t>
+                       ) -> bool {
+        method::lookup(fcx, op_ex, args[0], callee_id, opname,
+                       self_t, [], deref_args, CheckTraitsOnly,
+                       autoderef_receiver).is_some()
+    }
+
     fn lookup_op_method(fcx: @FnCtxt,
                         callee_id: ast::NodeId,
                         op_ex: &ast::Expr,
@@ -3195,27 +3210,36 @@ pub fn check_expr_with_unifier(fcx: @FnCtxt,
                       let resolved = structurally_resolved_type(fcx,
                                                                 expr.span,
                                                                 raw_base_t);
-                      let index_ident = tcx.sess.ident_of("index");
-                      let error_message = || {
-                        fcx.type_error_message(expr.span,
-                                               |actual| {
-                                                format!("cannot index a value \
-                                                      of type `{}`",
-                                                     actual)
-                                               },
-                                               base_t,
-                                               None);
+
+                      let mutable = match ty::get(resolved).sty {
+                        ty::ty_rptr(_, ty::mt{ ty: _, mutbl: ast::MutMutable }) => true,
+                        _ => false
                       };
-                      let ret_ty = lookup_op_method(fcx,
-                                                    callee_id,
-                                                    expr,
-                                                    resolved,
-                                                    index_ident.name,
-                                                    [base, idx],
-                                                    DoDerefArgs,
-                                                    AutoderefReceiver,
-                                                    error_message,
-                                                    expected);
+
+                      let lookup = |name| {
+                          let ident = tcx.sess.ident_of(name);
+                          lookup_op_method(fcx, callee_id, expr, resolved, ident.name,
+                                          [base, idx], DoDerefArgs, AutoderefReceiver,
+                                          || fcx.type_error_message(expr.span, |actual| {
+                                              format!("cannot {:s} a value of type `{:s}`", name, actual)
+                                          }, base_t, None), expected)
+                      };
+
+                      let has_op = |name| {
+                          let ident = tcx.sess.ident_of(name);
+                          has_op_method(fcx, callee_id, expr, resolved, ident.name,
+                                          [base, idx], DoDerefArgs, AutoderefReceiver,
+                                          expected)
+                      };
+
+                      let ret_ty = if mutable {
+                          lookup("index_mut")
+                      } else if has_op("index_ref") {
+                          lookup("index_ref")
+                      } else {
+                          lookup("index")
+                      };
+
                       fcx.write_ty(id, ret_ty);
                   }
               }

src/libstd/ops.rs

@@ -460,8 +460,18 @@ pub trait Shr<RHS,Result> {
  * ```
  */
 #[lang="index"]
-pub trait Index<Index,Result> {
-    fn index(&self, index: &Index) -> Result;
+pub trait Index<Element,Result> {
+    fn index(&self, element: &Element) -> Result;
+}
+
+#[lang="index_ref"]
+pub trait IndexRef<Element,Result> {
+    fn index<'a>(&'a self, element: &Element) -> &'a Result;
+}
+
+#[lang="index_mut"]
+pub trait IndexMut<Element,Result> {
+    fn index_mut<'a>(&'a mut self, element: &Element) -> &'a mut Result;
 }
 
 #[cfg(test)]

src/libstd/prelude.rs

@@ -33,7 +33,7 @@ pub use kinds::{Freeze, Pod, Send, Sized};
 pub use ops::{Add, Sub, Mul, Div, Rem, Neg, Not};
 pub use ops::{BitAnd, BitOr, BitXor};
 pub use ops::{Drop};
-pub use ops::{Shl, Shr, Index};
+pub use ops::{Shl, Shr, Index, IndexMut, IndexRef};
 pub use option::{Option, Some, None};
 pub use result::{Result, Ok, Err};
 

src/test/compile-fail/index_and_index_ref.rs

@@ -0,0 +1,35 @@
+// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+struct Foo {
+    bar: Bar
+}
+
+#[deriving(Clone)]
+struct Bar {
+    age: int
+}
+
+impl Index<int,Bar> for Foo {
+    fn index(&self, element: &int) -> Bar {
+        self.bar.clone()
+    }
+}
+
+impl IndexRef<int,Bar> for Foo {
+    fn index<'a>(&'a self, element: &int) -> &'a Bar {
+        &self.bar
+    }
+}
+
+fn main() {
+  let foo = Foo{ bar: Bar{ age: 50 } };
+  foo[10]; //~ ERROR multiple applicable methods in scope
+}

src/test/run-pass/overload-index-operator.rs

@@ -17,6 +17,27 @@ struct AssociationList<K,V> {
     pairs: ~[AssociationPair<K,V>]
 }
 
+#[deriving(Clone)]
+struct MyVec<T> {
+    values: ~[T]
+}
+
+impl<T> MyVec<T> {
+    pub fn new() -> MyVec<T> {
+      MyVec{ values: ~[] }
+    }
+
+    pub fn push(&mut self, val: T) {
+        self.values.push(val);
+    }
+}
+
+impl<T> IndexRef<int, T> for MyVec<T> {
+    fn index<'a>(&'a self, index: &int) -> &'a T {
+        &self.values[*index]
+    }
+}
+
 #[deriving(Clone)]
 struct AssociationPair<K,V> {
     key: K,
@@ -40,6 +61,17 @@ impl<K:Eq,V:Clone> Index<K,V> for AssociationList<K,V> {
     }
 }
 
+impl<K:Eq,V:Clone> IndexMut<K,V> for AssociationList<K,V> {
+    fn index_mut<'a>(&'a mut self, index: &K) -> &'a mut V {
+        for pair in self.pairs.mut_iter() {
+            if pair.key == *index {
+                &mut pair.value;
+            }
+        }
+        fail!("No value found for key: {:?}", index);
+    }
+}
+
 pub fn main() {
     let foo = ~"foo";
     let bar = ~"bar";
@@ -53,4 +85,19 @@ pub fn main() {
 
     assert!(list[foo] == 22)
     assert!(list[bar] == 44)
+
+    let mut list = AssociationList { pairs: ~[] };
+    list.push(foo.clone(), ~[1,2,3]);
+    list.push(bar.clone(), ~[5,6,7]);
+
+    assert_eq!(list[foo], ~[1,2,3]);
+    assert_eq!(list[bar], ~[5,6,7]);
+
+    let mut list = AssociationList{ pairs: ~[] };
+    list.push(foo.clone(), MyVec::new());
+    let mut my_vec = list[foo];
+    my_vec.push(1);
+    my_vec.push(2);
+    assert_eq!(my_vec[0], &1);
+    assert_eq!(my_vec[1], &2);
 }