expand memoryrefnew capabilities

Author: oscardssmith

Compiler/src/tfuncs.jl

@@ -2200,7 +2200,7 @@ function memoryref_builtin_common_nothrow(argtypes::Vector{Any})
         idx = widenconst(argtypes[2])
         idx ⊑ Int || return false
         boundscheck ⊑ Bool || return false
-        memtype ⊑ GenericMemoryRef || return false
+        memtype ⊑ Union{GenericMemory, GenericMemoryRef} || return false
         # If we have @inbounds (last argument is false), we're allowed to assume
         # we don't throw bounds errors.
         if isa(boundscheck, Const)

base/boot.jl

@@ -591,7 +591,7 @@ const undef = UndefInitializer()
 (self::Type{GenericMemory{kind,T,addrspace}})() where {T,kind,addrspace} = self(undef, 0)
 
 memoryref(mem::GenericMemory) = memoryrefnew(mem)
-memoryref(mem::GenericMemory, i::Integer) = memoryrefnew(memoryrefnew(mem), Int(i), @_boundscheck)
+memoryref(mem::GenericMemory, i::Integer) = memoryrefnew(me), Int(i), @_boundscheck)
 memoryref(ref::GenericMemoryRef, i::Integer) = memoryrefnew(ref, Int(i), @_boundscheck)
 GenericMemoryRef(mem::GenericMemory) = memoryref(mem)
 GenericMemoryRef(mem::GenericMemory, i::Integer) = memoryref(mem, i)

base/essentials.jl

@@ -383,7 +383,7 @@ default_access_order(a::GenericMemoryRef{:not_atomic}) = :not_atomic
 default_access_order(a::GenericMemoryRef{:atomic}) = :monotonic
 
 getindex(A::GenericMemory, i::Int) = (@_noub_if_noinbounds_meta;
-    memoryrefget(memoryrefnew(memoryrefnew(A), i, @_boundscheck), default_access_order(A), false))
+    memoryrefget(memoryrefnew(A, i, @_boundscheck), default_access_order(A), false))
 getindex(A::GenericMemoryRef) = memoryrefget(A, default_access_order(A), @_boundscheck)
 
 """
@@ -959,7 +959,7 @@ function setindex!(A::Array{Any}, @nospecialize(x), i::Int)
     memoryrefset!(memoryrefnew(getfield(A, :ref), i, false), x, :not_atomic, false)
     return A
 end
-setindex!(A::Memory{Any}, @nospecialize(x), i::Int) = (memoryrefset!(memoryrefnew(memoryrefnew(A), i, @_boundscheck), x, :not_atomic, @_boundscheck); A)
+setindex!(A::Memory{Any}, @nospecialize(x), i::Int) = (memoryrefset!(memoryrefnew(A, i, @_boundscheck), x, :not_atomic, @_boundscheck); A)
 setindex!(A::MemoryRef{T}, x) where {T} = (memoryrefset!(A, convert(T, x), :not_atomic, @_boundscheck); A)
 setindex!(A::MemoryRef{Any}, @nospecialize(x)) = (memoryrefset!(A, x, :not_atomic, @_boundscheck); A)
 

src/builtins.c

@@ -1765,30 +1765,44 @@ JL_CALLABLE(jl_f_memoryrefnew)
         return (jl_value_t*)jl_new_memoryref(typ, m, m->ptr);
     }
     else {
-        JL_TYPECHK(memoryrefnew, genericmemoryref, args[0]);
         JL_TYPECHK(memoryrefnew, long, args[1]);
         if (nargs == 3)
             JL_TYPECHK(memoryrefnew, bool, args[2]);
-        jl_genericmemoryref_t *m = (jl_genericmemoryref_t*)args[0];
         size_t i = jl_unbox_long(args[1]) - 1;
-        const jl_datatype_layout_t *layout = ((jl_datatype_t*)jl_typetagof(m->mem))->layout;
-        char *data = (char*)m->ptr_or_offset;
+        char *data;
+        if(jl_is_genericmemory(args[0])) {
+            jl_genericmemory_t *m = (jl_genericmemory_t*)args[0];
+            jl_value_t *typ = jl_apply_type((jl_value_t*)jl_genericmemoryref_type, jl_svec_data(((jl_datatype_t*)jl_typetagof(m))->parameters), 3);
+            if (i >= m->length)
+                jl_bounds_error((jl_value_t*)m, args[1]);
+            const jl_datatype_layout_t *layout = ((jl_datatype_t*)jl_typetagof(m))->layout;
+            if (layout->flags.arrayelem_isunion || layout->size == 0)
+                return (jl_value_t*)jl_new_memoryref(typ, m, (char*)i);
+            else if (layout->flags.arrayelem_isboxed)
+                return (jl_value_t*)jl_new_memoryref(typ, m, (char*)m->ptr + sizeof(jl_value_t*)*i);
+            return (jl_value_t*)jl_new_memoryref(typ, m, (char*)m->ptr + layout->size*i);
+        }
+        JL_TYPECHK(memoryrefnew, genericmemoryref, args[0]);
+        jl_genericmemoryref_t *m = (jl_genericmemoryref_t*)args[0];
+        jl_genericmemory_t *mem = m->mem;
+        data = (char*)m->ptr_or_offset;
+        const jl_datatype_layout_t *layout = ((jl_datatype_t*)jl_typetagof(mem))->layout;
         if (layout->flags.arrayelem_isboxed) {
-            if (((data - (char*)m->mem->ptr) / sizeof(jl_value_t*)) + i >= m->mem->length)
+            if (((data - (char*)mem->ptr) / sizeof(jl_value_t*)) + i >= mem->length)
                 jl_bounds_error((jl_value_t*)m, args[1]);
             data += sizeof(jl_value_t*) * i;
         }
         else if (layout->flags.arrayelem_isunion || layout->size == 0) {
-            if ((size_t)data + i >= m->mem->length)
+            if ((size_t)data + i >= mem->length)
                 jl_bounds_error((jl_value_t*)m, args[1]);
             data += i;
         }
         else {
-            if (((data - (char*)m->mem->ptr) / layout->size) + i >= m->mem->length)
+            if (((data - (char*)mem->ptr) / layout->size) + i >= mem->length)
                 jl_bounds_error((jl_value_t*)m, args[1]);
             data += layout->size * i;
         }
-        return (jl_value_t*)jl_new_memoryref((jl_value_t*)jl_typetagof(m), m->mem, data);
+        return (jl_value_t*)jl_new_memoryref((jl_value_t*)jl_typetagof(m), mem, data);
     }
 }
 

src/codegen.cpp

@@ -4138,16 +4138,25 @@ static bool emit_builtin_call(jl_codectx_t &ctx, jl_cgval_t *ret, jl_value_t *f,
 
     else if (f == BUILTIN(memoryrefnew) && (nargs == 2 || nargs == 3)) {
         const jl_cgval_t &ref = argv[1];
-        jl_value_t *mty_dt = jl_unwrap_unionall(ref.typ);
-        if (jl_is_genericmemoryref_type(mty_dt) && jl_is_concrete_type(mty_dt)) {
-            mty_dt = jl_field_type_concrete((jl_datatype_t*)mty_dt, 1);
-            const jl_datatype_layout_t *layout = ((jl_datatype_t*)mty_dt)->layout;
+        jl_datatype_t *mty_dt = (jl_datatype_t*)jl_unwrap_unionall(ref.typ);
+        if (jl_is_genericmemoryref_type(mty_dt) && jl_is_concrete_type((jl_value_t*)mty_dt)) {
+            mty_dt = (jl_datatype_t*)jl_field_type_concrete(mty_dt, 1);
+            const jl_datatype_layout_t *layout = mty_dt->layout;
             jl_value_t *boundscheck = nargs == 3 ? argv[3].constant : nullptr;
             if (nargs == 3)
                 emit_typecheck(ctx, argv[3], (jl_value_t*)jl_bool_type, "memoryrefnew");
             *ret = emit_memoryref(ctx, ref, argv[2], boundscheck, layout);
             return true;
         }
+        if (jl_is_genericmemory_type(mty_dt) && jl_is_concrete_type((jl_value_t*)mty_dt)) {
+            const jl_datatype_layout_t *layout = mty_dt->layout;
+            jl_value_t *boundscheck = nargs == 3 ? argv[3].constant : nullptr;
+            if (nargs == 3)
+                emit_typecheck(ctx, argv[3], (jl_value_t*)jl_bool_type, "memoryrefnew");
+            jl_value_t *typ = jl_apply_type((jl_value_t*)jl_genericmemoryref_type, jl_svec_data(mty_dt->parameters), jl_svec_len(mty_dt->parameters));
+            *ret = emit_memoryref(ctx, _emit_memoryref(ctx, ref, layout, typ), argv[2], boundscheck, layout);
+            return true;
+        }
     }
 
     else if (f == BUILTIN(memoryrefoffset) && nargs == 1) {