making size(A, k) give 1 for out-of-range dimensions

resolves issue #306
JuliaLang · Feb 10, 2012 · db7df14 · db7df14
1 parent 078d3a5
commit db7df14
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 7 deletions.
diff --git a/j/abstractarray.j b/j/abstractarray.j
@@ -12,7 +12,7 @@ typealias RangeIndex Union(Int, Range{Int}, Range1{Int})
 
 ## Basic functions ##
 
-size(t::AbstractArray, d) = size(t)[d]
+size{T,n}(t::AbstractArray{T,n}, d) = (d>n ? 1 : size(t)[d])
 eltype{T,n}(::AbstractArray{T,n}) = T
 ndims{T,n}(::AbstractArray{T,n}) = n
 numel(t::AbstractArray) = prod(size(t))
@@ -22,6 +22,9 @@ last(a::AbstractArray) = a[end]
 
 function stride(a::AbstractArray, i::Integer)
     s = 1
+    if i > ndims(a)
+        return numel(a)
+    end
     for n=1:(i-1)
         s *= size(a, n)
     end

diff --git a/j/string.j b/j/string.j
@@ -30,7 +30,7 @@ show(s::String) = print_quoted(s)
 
 size(s::String) = (length(s),)
 size(s::String, d::Integer) = d==1 ? length(s) :
-    error("in size: tupleref: index ",d," out of range")
+    error("in size: dimension ",d," out of range")
 
 strlen(s::DirectIndexString) = length(s)
 function strlen(s::String)

diff --git a/src/array.c b/src/array.c
@@ -207,8 +207,10 @@ JL_CALLABLE(jl_f_arraysize)
     if (nargs == 2) {
         JL_TYPECHK(arraysize, long, args[1]);
         int dno = jl_unbox_long(args[1]);
-        if (dno < 1 || dno > nd)
+        if (dno < 1)
             jl_error("arraysize: dimension out of range");
+        if (dno > nd)
+            return jl_box_long(1);
         return jl_box_long((&a->nrows)[dno-1]);
     }
     else {

diff --git a/src/codegen.cpp b/src/codegen.cpp
@@ -661,15 +661,37 @@ static Value *emit_known_call(jl_value_t *ff, jl_value_t **args, size_t nargs,
                         JL_GC_POP();
                         return emit_arraysize(ary, idx);
                     }
+                    else if (idx > ndims) {
+                        JL_GC_POP();
+                        return ConstantInt::get(T_size, 1);
+                    }
                 }
                 else {
                     Value *idx = emit_unbox(T_size, T_psize,
                                             emit_unboxed(args[2], ctx));
-                    emit_bounds_check(idx, ConstantInt::get(T_size,ndims),
-                                      "arraysize: dimension out of range",
-                                      ctx);
+                    error_unless(builder.CreateICmpSGT(idx,
+                                                      ConstantInt::get(T_size,0)),
+                                 "arraysize: dimension out of range", ctx);
+                    BasicBlock *outBB = BasicBlock::Create(getGlobalContext(),"outofrange",ctx->f);
+                    BasicBlock *inBB = BasicBlock::Create(getGlobalContext(),"inrange");
+                    BasicBlock *ansBB = BasicBlock::Create(getGlobalContext(),"arraysize");
+                    builder.CreateCondBr(builder.CreateICmpSLE(idx,
+                                                              ConstantInt::get(T_size, ndims)),
+                                         inBB, outBB);
+                    builder.SetInsertPoint(outBB);
+                    Value *v_one = ConstantInt::get(T_size, 1);
+                    builder.CreateBr(ansBB);
+                    ctx->f->getBasicBlockList().push_back(inBB);
+                    builder.SetInsertPoint(inBB);
+                    Value *v_sz = emit_arraysize(ary, idx);
+                    builder.CreateBr(ansBB);
+                    ctx->f->getBasicBlockList().push_back(ansBB);
+                    builder.SetInsertPoint(ansBB);
+                    PHINode *result = builder.CreatePHI(T_size, 2);
+                    result->addIncoming(v_one, outBB);
+                    result->addIncoming(v_sz, inBB);
                     JL_GC_POP();
-                    return emit_arraysize(ary, idx);
+                    return result;
                 }
             }
         }