Reference implementation and commentary for Group length extension

2 files changed, 10 insertions, 5 deletions

diff --git a/spec/primitive.md b/spec/primitive.md
index fca283d5..2662f8c4 100644
--- a/spec/primitive.md
+++ b/spec/primitive.md
@@ -161,7 +161,7 @@ For **Select** (`⊏`), `𝕨` is an array of natural numbers, or a list of such
 
 **First Cell** (`⊏`) selects the initial major cell of `𝕩`, giving an error if `𝕩` has rank 0 or length 0.
 
-**Group** (`⊔`) performs an opposite operation to Select, so that `𝕨` specifies not the argument index that result values come from, but the result index that argument values go to. The general case is that `𝕨` is a list of arrays of numbers; if it has depth less than 2 it's converted to this form by first enclosing it if it's an atom, then placing it in a length-1 list. After this transformation, the result rank is `≠𝕨`, and each result element has rank `(≠𝕨)+(=𝕩)-+´=¨𝕨`, with the initial `≠𝕨` axes corresponding to elements of `𝕨` and the remainder to trailing axes of `𝕩`. Each atom in `𝕨` can be either a natural number or `¯1` (which indicates the corresponding position in `𝕩` will be omitted). If `¯1` doesn't appear, the result has the property that each cell of `𝕩` appears in the corresponding element of `𝕨⊏𝕨⊔𝕩`. More concretely, the length of the result along axis `a` is the maximum value in `a⊑𝕨` plus one, or zero if `a⊑𝕨` is empty. Axis `a` corresponds to `=a⊑𝕨` axes in `𝕩`, and an element of the result at position `i` along this axis contains all positions in `𝕩` where `i=a⊑𝕨`. There may be multiple such positions, and they're arranged along axis `a` of that result element according to their index order in `𝕩`. **Group Indices** treats its argument `𝕩` as a left argument for Group and uses a right argument made up of indices, with the definition `⊔⟜(↕≠⚇1)`.
+**Group** (`⊔`) performs an opposite operation to Select, so that `𝕨` specifies not the argument index that result values come from, but the result index that argument values go to. The general case is that `𝕨` is a list of arrays of numbers; if it has depth less than 2 it's converted to this form by first enclosing it if it's an atom, then placing it in a length-1 list. After this transformation, the result rank is `≠𝕨`, and each result element has rank `(≠𝕨)+(=𝕩)-+´=¨𝕨`, with the initial `≠𝕨` axes corresponding to elements of `𝕨` and the remainder to trailing axes of `𝕩`. Each atom in `𝕨` can be either a natural number or `¯1` (which indicates the corresponding position in `𝕩` will be omitted). If `¯1` doesn't appear, the result has the property that each cell of `𝕩` appears in the corresponding element of `𝕨⊏𝕨⊔𝕩`. More concretely, the length of the result along axis `a` is the maximum value in `a⊑𝕨` plus one, or zero if `a⊑𝕨` is empty. Axis `a` corresponds to `=a⊑𝕨` axes in `𝕩`, and an element of the result at position `i` along this axis contains all positions in `𝕩` where `i=a⊑𝕨`. There may be multiple such positions, and they're arranged along axis `a` of that result element according to their index order in `𝕩`. The shapes of components of `𝕨` must match the corresponding axes of `𝕩`, except for rank-1 components of `𝕨`, which can match or have an extra element. This element, which like the others is either a natural number or `¯1`, gives the minimum length of the result axis corresponding to the component of `𝕨` in question, but otherwise does not affect the result. **Group Indices** treats its argument `𝕩` as a left argument for Group and uses a right argument made up of indices, with the definition `⊔⟜(↕≠⚇1)`.
 
 **Range** (`↕`) is extended to apply to a list of natural numbers, in addition to the provided case of a single natural number (an enclosed natural number `𝕩` should still result in an error). For a list `𝕩`, the result is an array of shape `𝕩` in which the value at a given index is that index, as a list of natural numbers. That is, `i≡i⊑↕𝕩` for any list of natural numbers `i` with `∧´i<𝕩`.
 
diff --git a/spec/reference.bqn b/spec/reference.bqn
index 288fdc39..988bfb51 100644
--- a/spec/reference.bqn
+++ b/spec/reference.bqn
@@ -383,11 +383,16 @@ Group←{
   𝕨↩Pair∘ToArray⍟(2>≡)𝕨
   ! 1==𝕨
   {!∧´Int¨𝕩⋄!∧´¯1≤𝕩}∘⥊¨𝕨
-  n←+´=¨𝕨
+  n←+´r←=¨𝕨
   ! n≤=𝕩
-  ! (∾≢⌜𝕨)≡n↑≢𝕩
-  𝕨↩⥊¨𝕨 ⋄ 𝕩↩((≠¨𝕨)∾n↓≢𝕩)⥊𝕩
-  (𝕨⊸=/𝕩˙)¨↕1+¯1⌈´¨𝕨
+  ld←(∾≢⌜𝕨)-n↑≢𝕩
+  ! ∧´(0⊸≤∧≤⟜(r/1=r))ld
+  dr←r⌊(0»+`r)⊏ld∾⟨0⟩
+  s←dr⊣◶⟨0,¯1⊸⊑⟩¨𝕨
+  𝕨↩dr(⥊¯1⊸↓⍟⊣)¨𝕨
+  s⌈↩1+¯1⌈´¨𝕨
+  𝕩↩((≠¨𝕨)∾n↓≢𝕩)⥊𝕩
+  (𝕨⊸=/𝕩˙)¨↕s
 }
 GroupInds←{
   ! 1==𝕩