Short-circuiting even when there is an offset

[athas]

I would like short circuiting to work for this program:

types {
  
}



entry("main",
      {ns: *[]i32},
      {*[]i32})
  entry_main (ns_mem_5212 : mem@device,
              ns_offset_5211 : i64,
              n_5168 : i64,
              ns_5169 : *[n_5168]i32 @ ns_mem_5212 ->
                        {offset: ns_offset_5211;
                         strides: [1i64];
                         shape: [n_5168]})
  : {mem@device,
     i64,
     *[n_5168]i32 @ ?0@device ->
     {offset: ?1;
      strides: [1i64];
      shape: [n_5168]}} = {
  let {segmap_tblock_size_5204 : i64} =
    get_size(segmap_tblock_size_5196, thread_block_size)
  let {segmap_usable_groups_5205 : i64} =
    sdiv_up64(n_5168, segmap_tblock_size_5204)
  let {bytes_5214 : i64} =
    mul_nw64(4i64, n_5168)
  let {mem_5215 : mem@device} =
    alloc(bytes_5214, @device)
  let {defunc_0_map_res_5206 : [n_5168]i32 @ mem_5215 ->
                               {offset: 0i64;
                                strides: [1i64];
                                shape: [n_5168]}} =
    segmap(thread; ; grid=segmap_usable_groups_5205; blocksize=segmap_tblock_size_5204)
    (gtid_5207 < n_5168) (~phys_tid_5208) : {i32} {
      let {eta_p_5209 : i32} =
        ns_5169[gtid_5207]
      let {lifted_lambda_res_5210 : i32} =
        add32(1i32, eta_p_5209)
      return {returns lifted_lambda_res_5210}
    }
  in {mem_5215, 0i64, defunc_0_map_res_5206}
}

Note that the array being indexed (ns_5169) has a nonzero (in fact unknown) offset. The compiler currently does not generate code like this, but I would like it to do so in the future, which means short circuiting would have to keep up. I don't think supporting this would be very difficult. It requires adjusting the LMAD of the result, but so does most other things that short circuiting performs.