Allow operator overloading of the indexing operator

The method `op_index` (which takes a single argument) is used for
this.

Issue #1520

Author: marijnh

src/comp/middle/resolve.rs

@@ -1983,9 +1983,12 @@ fn visit_mod_with_impl_scope(e: @env, m: ast::_mod, s: span, sc: iscopes,
 
 fn resolve_impl_in_expr(e: @env, x: @ast::expr, sc: iscopes, v: vt<iscopes>) {
     alt x.node {
+      // Store the visible impls in all exprs that might need them
       ast::expr_field(_, _, _) | ast::expr_path(_) | ast::expr_cast(_, _) |
       ast::expr_binary(_, _, _) | ast::expr_unary(_, _) |
-      ast::expr_assign_op(_, _, _) { e.impl_map.insert(x.id, sc); }
+      ast::expr_assign_op(_, _, _) | ast::expr_index(_, _) {
+        e.impl_map.insert(x.id, sc);
+      }
       _ {}
     }
     visit::visit_expr(x, sc, v);

src/comp/middle/trans.rs

@@ -2724,9 +2724,8 @@ fn trans_rec_field(bcx: @block_ctxt, base: @ast::expr,
     ret {bcx: bcx, val: val, kind: owned};
 }
 
-fn trans_index(cx: @block_ctxt, sp: span, base: @ast::expr, idx: @ast::expr,
-               id: ast::node_id) -> lval_result {
-    // Is this an interior vector?
+fn trans_index(cx: @block_ctxt, ex: @ast::expr, base: @ast::expr,
+               idx: @ast::expr) -> lval_result {
     let base_ty = ty::expr_ty(bcx_tcx(cx), base);
     let exp = trans_temp_expr(cx, base);
     let lv = autoderef(exp.bcx, exp.val, base_ty);
@@ -2745,7 +2744,7 @@ fn trans_index(cx: @block_ctxt, sp: span, base: @ast::expr, idx: @ast::expr,
         ix_val = Trunc(bcx, ix.val, ccx.int_type);
     } else { ix_val = ix.val; }
 
-    let unit_ty = node_id_type(bcx_ccx(cx), id);
+    let unit_ty = node_id_type(bcx_ccx(cx), ex.id);
     let unit_sz = size_of(bcx, unit_ty);
     bcx = unit_sz.bcx;
     maybe_name_value(bcx_ccx(cx), unit_sz.val, "unit_sz");
@@ -2760,11 +2759,11 @@ fn trans_index(cx: @block_ctxt, sp: span, base: @ast::expr, idx: @ast::expr,
     CondBr(bcx, bounds_check, next_cx.llbb, fail_cx.llbb);
     // fail: bad bounds check.
 
-    trans_fail(fail_cx, some::<span>(sp), "bounds check");
+    trans_fail(fail_cx, some(ex.span), "bounds check");
     let elt =
         if check type_has_static_size(ncx, unit_ty) {
             let elt_1 = GEP(next_cx, body, [ix_val]);
-            let llunitty = type_of(ncx, sp, unit_ty);
+            let llunitty = type_of(ncx, ex.span, unit_ty);
             PointerCast(next_cx, elt_1, T_ptr(llunitty))
         } else {
             body = PointerCast(next_cx, body, T_ptr(T_i8()));
@@ -2812,7 +2811,7 @@ fn trans_lval(cx: @block_ctxt, e: @ast::expr) -> lval_result {
         ret trans_rec_field(cx, base, ident);
       }
       ast::expr_index(base, idx) {
-        ret trans_index(cx, e.span, base, idx, e.id);
+        ret trans_index(cx, e, base, idx);
       }
       ast::expr_unary(ast::deref, base) {
         let ccx = bcx_ccx(cx);
@@ -3560,6 +3559,15 @@ fn trans_expr(bcx: @block_ctxt, e: @ast::expr, dest: dest) -> @block_ctxt {
       ast::expr_field(_, _, _) {
         fail "Taking the value of a method does not work yet (issue #435)";
       }
+      ast::expr_index(base, idx) {
+        // If it is here, it's not an lval, so this is a user-defined index op
+        let origin = bcx_ccx(bcx).method_map.get(e.id);
+        let callee_id = ast_util::op_expr_callee_id(e);
+        let fty = ty::node_id_to_monotype(tcx, callee_id);
+        ret trans_call_inner(bcx, fty, {|bcx|
+            trans_impl::trans_method_callee(bcx, callee_id, base, origin)
+        }, [idx], e.id, dest);
+      }
 
       // These return nothing
       ast::expr_break {

src/comp/middle/ty.rs

@@ -1607,9 +1607,10 @@ fn expr_has_ty_params(cx: ctxt, expr: @ast::expr) -> bool {
 
 fn expr_is_lval(method_map: typeck::method_map, e: @ast::expr) -> bool {
     alt e.node {
-      ast::expr_path(_) | ast::expr_index(_, _) |
-      ast::expr_unary(ast::deref, _) { true }
-      ast::expr_field(base, ident, _) { !method_map.contains_key(e.id) }
+      ast::expr_path(_) | ast::expr_unary(ast::deref, _) { true }
+      ast::expr_field(_, _, _) | ast::expr_index(_, _) {
+        !method_map.contains_key(e.id)
+      }
       _ { false }
     }
 }

src/comp/middle/typeck.rs

@@ -1780,6 +1780,22 @@ fn check_expr_with_unifier(fcx: @fn_ctxt, expr: @ast::expr, unify: unifier,
           _ { none }
         }
     }
+    fn lookup_op_method(fcx: @fn_ctxt, op_ex: @ast::expr, self_t: ty::t,
+                        opname: str,
+                        args: [option::t<@ast::expr>]) -> option::t<ty::t> {
+        let isc = fcx.ccx.impl_map.get(op_ex.id);
+        alt lookup_method(fcx, isc, opname, self_t, op_ex.span) {
+          some({method_ty, n_tps: 0u, substs, origin}) {
+            let callee_id = ast_util::op_expr_callee_id(op_ex);
+            write::ty_fixup(fcx, callee_id, {substs: some(substs),
+                                             ty: method_ty});
+            check_call_or_bind(fcx, op_ex.span, method_ty, args);
+            fcx.ccx.method_map.insert(op_ex.id, origin);
+            some(ty::ty_fn_ret(fcx.ccx.tcx, method_ty))
+          }
+          _ { none }
+        }
+    }
     fn check_binop(fcx: @fn_ctxt, ex: @ast::expr, ty: ty::t,
                    op: ast::binop, rhs: @ast::expr) -> ty::t {
         let resolved_t = structurally_resolved_type(fcx, ex.span, ty);
@@ -1792,22 +1808,14 @@ fn check_expr_with_unifier(fcx: @fn_ctxt, expr: @ast::expr, unify: unifier,
             };
         }
 
-        let isc = fcx.ccx.impl_map.get(ex.id);
         alt binop_method(op) {
           some(name) {
-            alt lookup_method(fcx, isc, name, resolved_t, ex.span) {
-              some({method_ty, n_tps: 0u, substs, origin}) {
-                let callee_id = ast_util::op_expr_callee_id(ex);
-                write::ty_fixup(fcx, callee_id, {substs: some(substs),
-                                                 ty: method_ty});
-                check_call_or_bind(fcx, ex.span, method_ty, [some(rhs)]);
-                fcx.ccx.method_map.insert(ex.id, origin);
-                ret ty::ty_fn_ret(tcx, method_ty);
-              }
+            alt lookup_op_method(fcx, ex, resolved_t, name, [some(rhs)]) {
+              some(ret_ty) { ret ret_ty; }
               _ {}
             }
           }
-          none {}
+          _ {}
         }
         tcx.sess.span_err(
             ex.span, "binary operation " + ast_util::binop_to_str(op) +
@@ -1817,23 +1825,15 @@ fn check_expr_with_unifier(fcx: @fn_ctxt, expr: @ast::expr, unify: unifier,
     }
     fn check_user_unop(fcx: @fn_ctxt, op_str: str, mname: str,
                        ex: @ast::expr, rhs_t: ty::t) -> ty::t {
-        let isc = fcx.ccx.impl_map.get(ex.id);
-        let tcx = fcx.ccx.tcx;
-        alt lookup_method(fcx, isc, mname, rhs_t, ex.span) {
-          some({method_ty, n_tps: 0u, substs, origin}) {
-            let callee_id = ast_util::op_expr_callee_id(ex);
-            write::ty_fixup(fcx, callee_id, {substs: some(substs),
-                                             ty: method_ty});
-            check_call_or_bind(fcx, ex.span, method_ty, []);
-            fcx.ccx.method_map.insert(ex.id, origin);
-            ret ty::ty_fn_ret(tcx, method_ty);
+        alt lookup_op_method(fcx, ex, rhs_t, mname, []) {
+          some(ret_ty) { ret_ty }
+          _ {
+            fcx.ccx.tcx.sess.span_err(
+                ex.span, #fmt["cannot apply unary operator `%s` to type `%s`",
+                              op_str, ty_to_str(fcx.ccx.tcx, rhs_t)]);
+            rhs_t
           }
-          _ {}
         }
-        tcx.sess.span_err(
-            ex.span, #fmt["can not apply unary operator `%s` to type `%s`",
-                          op_str, ty_to_str(tcx, rhs_t)]);
-        rhs_t
     }
 
     let tcx = fcx.ccx.tcx;
@@ -1888,7 +1888,7 @@ fn check_expr_with_unifier(fcx: @fn_ctxt, expr: @ast::expr, unify: unifier,
               }
               ty::ty_ptr(inner) {
                 oper_t = inner.ty;
-                require_unsafe(fcx.ccx.tcx.sess, fcx.purity, expr.span);
+                require_unsafe(tcx.sess, fcx.purity, expr.span);
               }
               _ {
                 tcx.sess.span_fatal(expr.span,
@@ -1989,7 +1989,7 @@ fn check_expr_with_unifier(fcx: @fn_ctxt, expr: @ast::expr, unify: unifier,
       ast::expr_copy(a) {
         bot = check_expr_with_unifier(fcx, a, unify, expected);
         let tpot =
-            ty::node_id_to_ty_param_substs_opt_and_ty(fcx.ccx.tcx, a.id);
+            ty::node_id_to_ty_param_substs_opt_and_ty(tcx, a.id);
         write::ty_fixup(fcx, id, tpot);
 
       }
@@ -2073,9 +2073,8 @@ fn check_expr_with_unifier(fcx: @fn_ctxt, expr: @ast::expr, unify: unifier,
         let proto = alt ty::struct(tcx, expected) {
           ty::ty_fn({proto, _}) { proto }
           _ {
-            fcx.ccx.tcx.sess.span_warn(
-                expr.span,
-                "unable to infer kind of closure, defaulting to block");
+            tcx.sess.span_warn(expr.span, "unable to infer kind of closure, \
+                                           defaulting to block");
             ast::proto_block
           }
         };
@@ -2190,10 +2189,10 @@ fn check_expr_with_unifier(fcx: @fn_ctxt, expr: @ast::expr, unify: unifier,
         vec::reserve(elt_ts, vec::len(elts));
         for e in elts {
             check_expr(fcx, e);
-            let ety = expr_ty(fcx.ccx.tcx, e);
+            let ety = expr_ty(tcx, e);
             elt_ts += [ety];
         }
-        let typ = ty::mk_tup(fcx.ccx.tcx, elt_ts);
+        let typ = ty::mk_tup(tcx, elt_ts);
         write::ty_only_fixup(fcx, id, typ);
       }
       ast::expr_rec(fields, base) {
@@ -2312,26 +2311,39 @@ fn check_expr_with_unifier(fcx: @fn_ctxt, expr: @ast::expr, unify: unifier,
       }
       ast::expr_index(base, idx) {
         bot |= check_expr(fcx, base);
-        let base_t = expr_ty(tcx, base);
-        base_t = do_autoderef(fcx, expr.span, base_t);
+        let raw_base_t = expr_ty(tcx, base);
+        let base_t = do_autoderef(fcx, expr.span, raw_base_t);
         bot |= check_expr(fcx, idx);
         let idx_t = expr_ty(tcx, idx);
-        if !type_is_integral(fcx, idx.span, idx_t) {
-            tcx.sess.span_err(idx.span,
-                              "mismatched types: expected \
-                               `integer` but found `"
-                                  + ty_to_str(tcx, idx_t) + "`");
+        fn require_integral(fcx: @fn_ctxt, sp: span, t: ty::t) {
+            if !type_is_integral(fcx, sp, t) {
+                fcx.ccx.tcx.sess.span_err(sp, "mismatched types: expected \
+                                               `integer` but found `"
+                                  + ty_to_str(fcx.ccx.tcx, t) + "`");
+            }
         }
         alt structure_of(fcx, expr.span, base_t) {
-          ty::ty_vec(mt) { write::ty_only_fixup(fcx, id, mt.ty); }
+          ty::ty_vec(mt) {
+            require_integral(fcx, idx.span, idx_t);
+            write::ty_only_fixup(fcx, id, mt.ty);
+          }
           ty::ty_str {
+            require_integral(fcx, idx.span, idx_t);
             let typ = ty::mk_mach_uint(tcx, ast::ty_u8);
             write::ty_only_fixup(fcx, id, typ);
           }
           _ {
-            tcx.sess.span_fatal(expr.span,
-                                "vector-indexing bad type: " +
-                                    ty_to_str(tcx, base_t));
+            let resolved = structurally_resolved_type(fcx, expr.span,
+                                                      raw_base_t);
+            alt lookup_op_method(fcx, expr, resolved, "op_index",
+                                 [some(idx)]) {
+              some(ret_ty) { write::ty_only_fixup(fcx, id, ret_ty); }
+              _ {
+                tcx.sess.span_fatal(
+                    expr.span, "cannot index a value of type `" +
+                    ty_to_str(tcx, base_t) + "`");
+              }
+            }
           }
         }
       }
@@ -2922,7 +2934,8 @@ mod dict {
           }
           // Must resolve bounds on methods with bounded params
           ast::expr_field(_, _, _) | ast::expr_binary(_, _, _) |
-          ast::expr_unary(_, _) | ast::expr_assign_op(_, _, _) {
+          ast::expr_unary(_, _) | ast::expr_assign_op(_, _, _) |
+          ast::expr_index(_, _) {
             alt cx.method_map.find(ex.id) {
               some(method_static(did)) {
                 let bounds = ty::lookup_item_type(cx.tcx, did).bounds;

src/comp/syntax/parse/parser.rs

@@ -1008,6 +1008,7 @@ fn parse_dot_or_call_expr_with(p: parser, e0: pexpr) -> pexpr {
             let ix = parse_expr(p);
             hi = ix.span.hi;
             expect(p, token::RBRACKET);
+            p.get_id(); // see ast_util::op_expr_callee_id
             e = mk_pexpr(p, lo, hi, ast::expr_index(to_expr(e), ix));
           }
 

src/test/compile-fail/minus-string.rs

@@ -1,3 +1,3 @@
-// error-pattern:can not apply unary operator `-` to type `str`
+// error-pattern:cannot apply unary operator `-` to type `str`
 
 fn main() { -"foo"; }

src/test/run-pass/operator-overloading.rs

@@ -7,11 +7,16 @@ impl add_point for point {
     fn op_neg() -> point {
         {x: -self.x, y: -self.y}
     }
+    fn op_index(x: bool) -> int {
+        x ? self.x : self.y
+    }
 }
 
 fn main() {
     let p = {x: 10, y: 20};
     p += {x: 1, y: 2};
     assert p + {x: 5, y: 5} == {x: 16, y: 27};
     assert -p == {x: -11, y: -22};
+    assert p[true] == 11;
+    assert p[false] == 22;
 }