Show concisely arrays replicated from a single element (scalar or not)

1 files changed, 17 insertions, 3 deletions

diff --git a/src/Data/Array/Nested/Internal/Mixed.hs b/src/Data/Array/Nested/Internal/Mixed.hs
index 5730354..53b4f05 100644
--- a/src/Data/Array/Nested/Internal/Mixed.hs
+++ b/src/Data/Array/Nested/Internal/Mixed.hs
@@ -197,9 +197,17 @@ newtype ShowViaToListLinear sh a = ShowViaToListLinear (Mixed sh a)
 
 instance (Show a, Elt a) => Show (ShowViaToListLinear sh a) where
   showsPrec d (ShowViaToListLinear arr) = showParen (d > 10) $
-    -- TODO: to avoid ambiguity, this should type-apply the shape to mfromListLinear
-    showString "mfromListLinear " . shows (shxToList (mshape arr)) . showString " "
-      . shows (mtoListLinear arr)
+    let defaultResult =
+          -- TODO: to avoid ambiguity, this should type-apply the shape to mfromListLinear
+          showString "mfromListLinear " . shows (shxToList (mshape arr)) . showString " "
+            . shows (mtoListLinear arr)
+    in if stridesAreZero (shxLength $ mshape arr) (mstrideTree arr)
+       then case mtoListLinear arr of
+         [] -> defaultResult
+         [_] -> defaultResult
+         hd : _ -> showString "mreplicate " . shows (shxToList (mshape arr)) . showString " "
+                     . showsPrec 11 hd
+       else defaultResult
 
 newtype ShowViaPrimitive sh a = ShowViaPrimitive (Mixed sh (Primitive a))
 
@@ -279,6 +287,12 @@ data StrideTree =
     StrideLeaf [Int]
   | StrideNode StrideTree StrideTree
 
+stridesAreZero :: Int -> StrideTree -> Bool
+stridesAreZero prefixLen (StrideLeaf ss) =
+  all (== 0) (take prefixLen ss)
+stridesAreZero prefixLen (StrideNode ss1 ss2) =
+  stridesAreZero prefixLen ss1 && stridesAreZero prefixLen ss2
+
 -- | Allowable element types in a mixed array, and by extension in a 'Ranked' or
 -- 'Shaped' array. Note the polymorphic instance for 'Elt' of @'Primitive'
 -- a@; see the documentation for 'Primitive' for more details.