Reference implementation and commentary for Group length extension

3 files changed, 11 insertions, 6 deletions

diff --git a/docs/spec/primitive.html b/docs/spec/primitive.html
index ab65ef6c..d0eb9b2f 100644
--- a/docs/spec/primitive.html
+++ b/docs/spec/primitive.html
@@ -127,7 +127,7 @@
 <p><strong>First</strong> (<code><span class='Function'>⊑</span></code>) simply takes the first element of its argument in index order, or the fill element if <code><span class='Value'>𝕩</span></code> is empty.</p>
 <p>For <strong>Select</strong> (<code><span class='Function'>⊏</span></code>), <code><span class='Value'>𝕨</span></code> is an array of natural numbers, or a list of such arrays; if it's an empty list, it's interpreted as the former. The given arrays are matched with leading axes of <code><span class='Value'>𝕩</span></code> and used to select from those axes. Their shape is retained, so that the final shape is the combined shapes of each array of natural numbers in <code><span class='Value'>𝕨</span></code> in order, followed by the trailing (unmatched) shape of <code><span class='Value'>𝕩</span></code>. This means that a single axis in <code><span class='Value'>𝕩</span></code> can correspond to any number of axes in <code><span class='Value'>𝕨</span><span class='Function'>⊏</span><span class='Value'>𝕩</span></code>, depending on the rank of that portion of <code><span class='Value'>𝕨</span></code>. More precisely, the value of the result at an index <code><span class='Value'>j</span></code> is obtained by splitting <code><span class='Value'>j</span></code> into one index into each array of <code><span class='Value'>𝕨</span></code> followed by a partial index into <code><span class='Value'>𝕩</span></code>. An index <code><span class='Value'>i</span></code> for <code><span class='Value'>𝕩</span></code> comes from selecting from each array of <code><span class='Value'>𝕨</span></code> and appending the results to the partial index from <code><span class='Value'>j</span></code>, and the value <code><span class='Value'>i</span><span class='Function'>⊑</span><span class='Value'>𝕩</span></code> is <code><span class='Value'>j</span><span class='Function'>⊑</span><span class='Value'>𝕨</span><span class='Function'>⊏</span><span class='Value'>𝕩</span></code>.</p>
 <p><strong>First Cell</strong> (<code><span class='Function'>⊏</span></code>) selects the initial major cell of <code><span class='Value'>𝕩</span></code>, giving an error if <code><span class='Value'>𝕩</span></code> has rank 0 or length 0.</p>
-<p><strong>Group</strong> (<code><span class='Function'>⊔</span></code>) performs an opposite operation to Select, so that <code><span class='Value'>𝕨</span></code> specifies not the argument index that result values come from, but the result index that argument values go to. The general case is that <code><span class='Value'>𝕨</span></code> 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 <code><span class='Function'>≠</span><span class='Value'>𝕨</span></code>, and each result element has rank <code><span class='Paren'>(</span><span class='Function'>≠</span><span class='Value'>𝕨</span><span class='Paren'>)</span><span class='Function'>+</span><span class='Paren'>(</span><span class='Function'>=</span><span class='Value'>𝕩</span><span class='Paren'>)</span><span class='Function'>-+</span><span class='Modifier'>´</span><span class='Function'>=</span><span class='Modifier'>¨</span><span class='Value'>𝕨</span></code>, with the initial <code><span class='Function'>≠</span><span class='Value'>𝕨</span></code> axes corresponding to elements of <code><span class='Value'>𝕨</span></code> and the remainder to trailing axes of <code><span class='Value'>𝕩</span></code>. Each atom in <code><span class='Value'>𝕨</span></code> can be either a natural number or <code><span class='Number'>¯1</span></code> (which indicates the corresponding position in <code><span class='Value'>𝕩</span></code> will be omitted). If <code><span class='Number'>¯1</span></code> doesn't appear, the result has the property that each cell of <code><span class='Value'>𝕩</span></code> appears in the corresponding element of <code><span class='Value'>𝕨</span><span class='Function'>⊏</span><span class='Value'>𝕨</span><span class='Function'>⊔</span><span class='Value'>𝕩</span></code>. More concretely, the length of the result along axis <code><span class='Value'>a</span></code> is the maximum value in <code><span class='Value'>a</span><span class='Function'>⊑</span><span class='Value'>𝕨</span></code> plus one, or zero if <code><span class='Value'>a</span><span class='Function'>⊑</span><span class='Value'>𝕨</span></code> is empty. Axis <code><span class='Value'>a</span></code> corresponds to <code><span class='Function'>=</span><span class='Value'>a</span><span class='Function'>⊑</span><span class='Value'>𝕨</span></code> axes in <code><span class='Value'>𝕩</span></code>, and an element of the result at position <code><span class='Value'>i</span></code> along this axis contains all positions in <code><span class='Value'>𝕩</span></code> where <code><span class='Value'>i</span><span class='Function'>=</span><span class='Value'>a</span><span class='Function'>⊑</span><span class='Value'>𝕨</span></code>. There may be multiple such positions, and they're arranged along axis <code><span class='Value'>a</span></code> of that result element according to their index order in <code><span class='Value'>𝕩</span></code>. <strong>Group Indices</strong> treats its argument <code><span class='Value'>𝕩</span></code> as a left argument for Group and uses a right argument made up of indices, with the definition <code><span class='Function'>⊔</span><span class='Modifier2'>⟜</span><span class='Paren'>(</span><span class='Function'>↕≠</span><span class='Modifier2'>⚇</span><span class='Number'>1</span><span class='Paren'>)</span></code>.</p>
+<p><strong>Group</strong> (<code><span class='Function'>⊔</span></code>) performs an opposite operation to Select, so that <code><span class='Value'>𝕨</span></code> specifies not the argument index that result values come from, but the result index that argument values go to. The general case is that <code><span class='Value'>𝕨</span></code> 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 <code><span class='Function'>≠</span><span class='Value'>𝕨</span></code>, and each result element has rank <code><span class='Paren'>(</span><span class='Function'>≠</span><span class='Value'>𝕨</span><span class='Paren'>)</span><span class='Function'>+</span><span class='Paren'>(</span><span class='Function'>=</span><span class='Value'>𝕩</span><span class='Paren'>)</span><span class='Function'>-+</span><span class='Modifier'>´</span><span class='Function'>=</span><span class='Modifier'>¨</span><span class='Value'>𝕨</span></code>, with the initial <code><span class='Function'>≠</span><span class='Value'>𝕨</span></code> axes corresponding to elements of <code><span class='Value'>𝕨</span></code> and the remainder to trailing axes of <code><span class='Value'>𝕩</span></code>. Each atom in <code><span class='Value'>𝕨</span></code> can be either a natural number or <code><span class='Number'>¯1</span></code> (which indicates the corresponding position in <code><span class='Value'>𝕩</span></code> will be omitted). If <code><span class='Number'>¯1</span></code> doesn't appear, the result has the property that each cell of <code><span class='Value'>𝕩</span></code> appears in the corresponding element of <code><span class='Value'>𝕨</span><span class='Function'>⊏</span><span class='Value'>𝕨</span><span class='Function'>⊔</span><span class='Value'>𝕩</span></code>. More concretely, the length of the result along axis <code><span class='Value'>a</span></code> is the maximum value in <code><span class='Value'>a</span><span class='Function'>⊑</span><span class='Value'>𝕨</span></code> plus one, or zero if <code><span class='Value'>a</span><span class='Function'>⊑</span><span class='Value'>𝕨</span></code> is empty. Axis <code><span class='Value'>a</span></code> corresponds to <code><span class='Function'>=</span><span class='Value'>a</span><span class='Function'>⊑</span><span class='Value'>𝕨</span></code> axes in <code><span class='Value'>𝕩</span></code>, and an element of the result at position <code><span class='Value'>i</span></code> along this axis contains all positions in <code><span class='Value'>𝕩</span></code> where <code><span class='Value'>i</span><span class='Function'>=</span><span class='Value'>a</span><span class='Function'>⊑</span><span class='Value'>𝕨</span></code>. There may be multiple such positions, and they're arranged along axis <code><span class='Value'>a</span></code> of that result element according to their index order in <code><span class='Value'>𝕩</span></code>. The shapes of components of <code><span class='Value'>𝕨</span></code> must match the corresponding axes of <code><span class='Value'>𝕩</span></code>, except for rank-1 components of <code><span class='Value'>𝕨</span></code>, which can match or have an extra element. This element, which like the others is either a natural number or <code><span class='Number'>¯1</span></code>, gives the minimum length of the result axis corresponding to the component of <code><span class='Value'>𝕨</span></code> in question, but otherwise does not affect the result. <strong>Group Indices</strong> treats its argument <code><span class='Value'>𝕩</span></code> as a left argument for Group and uses a right argument made up of indices, with the definition <code><span class='Function'>⊔</span><span class='Modifier2'>⟜</span><span class='Paren'>(</span><span class='Function'>↕≠</span><span class='Modifier2'>⚇</span><span class='Number'>1</span><span class='Paren'>)</span></code>.</p>
 <p><strong>Range</strong> (<code><span class='Function'>↕</span></code>) 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 <code><span class='Value'>𝕩</span></code> should still result in an error). For a list <code><span class='Value'>𝕩</span></code>, the result is an array of shape <code><span class='Value'>𝕩</span></code> in which the value at a given index is that index, as a list of natural numbers. That is, <code><span class='Value'>i</span><span class='Function'>≡</span><span class='Value'>i</span><span class='Function'>⊑↕</span><span class='Value'>𝕩</span></code> for any list of natural numbers <code><span class='Value'>i</span></code> with <code><span class='Function'>∧</span><span class='Modifier'>´</span><span class='Value'>i</span><span class='Function'>&lt;</span><span class='Value'>𝕩</span></code>.</p>
 <p><strong>Indices</strong> (<code><span class='Function'>/</span></code>) applies to a list of natural numbers, and returns a list of natural numbers. The result contains <code><span class='Value'>i</span><span class='Function'>⊑</span><span class='Value'>𝕩</span></code> copies of each natural number index <code><span class='Value'>i</span></code> for <code><span class='Value'>𝕩</span></code>, in increasing order.</p>
 <h3 id="structural-manipulation">Structural manipulation</h3>
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==𝕩