Special Take/Drop implementation for atom 𝕨

1 files changed, 20 insertions, 7 deletions

diff --git a/src/r1.bqn b/src/r1.bqn
index f5a295ee..cd0d1c3f 100644
--- a/src/r1.bqn
+++ b/src/r1.bqn
@@ -432,13 +432,26 @@ structFn ← {
 }
 
 _takeDrop←{
-  ⟨gl,Noop,_inds⟩←𝕗
+  take ← 1 - 𝕗
+  gl   ← 𝕗⊑"↑"‿"↓"
+  noop ← 𝕗⊑⟨1-=⟜|, 1-0⊸=⟩
+  inds ← 𝕗⊑⟨
+    { 𝔽⍟(𝕨⊸<)a←|𝕩 ⋄ (0<𝕩)◶⟨¯∞⍟(<⟜0)⌜+⟜(𝕨+𝕩)⌜, ¯∞⍟(𝕨⊸≤)⌜⟩↕a }
+    { 𝔽 ⋄ 0⊸<◶⟨↕0⌈+,<∘⊢+⌜·↕0⌈-⟩ }
+  ⟩
   pre ← "𝕨"∾gl∾"𝕩: 𝕨 must "
   ernk ← pre∾"have rank at most 1"
   eint ← pre∾"consist of integers"
-  {
+  IsArray∘⊣◶{
+    eint ! Int 𝕨
+    p←0≤𝕨
+    l←𝕨p◶⟨0⌈+,⌊⟩≠𝕩
+    F←𝕩{(Fill𝕗)˙⌜↕𝕩}
+    k←1⋄S←⊢ ⋄ 𝕨⊸{k×´↩c←1 Cell𝕩 ⋄ S↩(⟨(0⌈(≠𝕩)-⊢)⍟(1-take)|𝕨⟩∾c)⊸⥊}⍟(1<=) 𝕩
+    S ((|∘𝕨-≠){𝕩p◶⟨∾˜,∾⟩F𝕨×k}⍟(>⟜0)⊢)⍟take (l×k) (take=p)◶↓‿↑ ⥊𝕩
+  }‿{
     ernk ! 1≥=𝕨
-    𝕨 ↩ Deshape 𝕨
+    𝕨 ↩ ⥊𝕨
     eint ! 1×´Int⌜𝕨
     r ← ≠𝕨
     s ← r {(1⌜∘↕𝕨-≠𝕩)∾𝕩}⍟(>⟜≠) ≢𝕩
@@ -452,11 +465,11 @@ _takeDrop←{
       Sel ← ⊑⟜(⥊𝕩)
       𝕩{Sel↩0⊸≤◶⟨(Fill𝕨)˙,Sel⟩}⍟⊢doFil
       Sel⌜ sh ⥊ i
-    }_fillBy_⊢ ToArray 𝕩
-  }
+    } 𝕩
+  }_fillBy_⊢ ⟜ ToArray
 }
-Take ← ⟨"↑" ⋄ 1-=⟜| ⋄ { 𝔽⍟(𝕨⊸<)a←|𝕩 ⋄ (0<𝕩)◶⟨¯∞⍟(<⟜0)⌜+⟜(𝕨+𝕩)⌜, ¯∞⍟(𝕨⊸≤)⌜⟩↕a }⟩_takeDrop
-Drop ← ⟨"↓" ⋄ 1-0⊸= ⋄ { 𝔽 ⋄ 0⊸<◶⟨↕0⌈+,<∘⊢+⌜·↕0⌈-⟩ }⟩_takeDrop
+Take ← 0 _takeDrop
+Drop ← 1 _takeDrop
 
 ShiftCheck←{
   "« or »: 𝕩 must have rank at least 1" ! 1≤=𝕩