From 2044be5a83d77536daed0f0f34d2baa6aa548dd1 Mon Sep 17 00:00:00 2001
From: Alberto Ruiz <aruiz@um.es>
Date: Wed, 5 Nov 2008 14:01:11 +0000
Subject: improved constant

---
 lib/Data/Packed/ST.hs | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

(limited to 'lib/Data/Packed/ST.hs')

diff --git a/lib/Data/Packed/ST.hs b/lib/Data/Packed/ST.hs
index 91b7968..9c9c3b3 100644
--- a/lib/Data/Packed/ST.hs
+++ b/lib/Data/Packed/ST.hs
@@ -1,4 +1,4 @@
-{-# OPTIONS -XTypeOperators -XRank2Types  -XFlexibleContexts #-}
+{-# OPTIONS -XTypeOperators -XRank2Types  -XFlexibleContexts -XBangPatterns #-}
 
 -----------------------------------------------------------------------------
 -- |
@@ -23,8 +23,10 @@ module Data.Packed.ST (
     STMatrix, newMatrix, thawMatrix, freezeMatrix, runSTMatrix,
     readMatrix, writeMatrix, modifyMatrix, liftSTMatrix,
     -- * Unsafe functions
+    newUndefinedVector,
     unsafeReadVector, unsafeWriteVector,
     unsafeThawVector, unsafeFreezeVector,
+    newUndefinedMatrix,
     unsafeReadMatrix, unsafeWriteMatrix,
     unsafeThawMatrix, unsafeFreezeMatrix
 ) where
@@ -87,9 +89,17 @@ readVector = safeIndexV unsafeReadVector
 writeVector :: Storable t => STVector s t -> Int -> t -> ST s ()
 writeVector = safeIndexV unsafeWriteVector
 
+{-# NOINLINE newUndefinedVector #-}
+newUndefinedVector :: Element t => Int -> ST s (STVector s t)
+newUndefinedVector = unsafeIOToST . fmap STVector . createVector
+
 {-# NOINLINE newVector #-}
 newVector :: Element t => t -> Int -> ST s (STVector s t)
-newVector v = unsafeThawVector . constant v
+newVector x n = do
+    v <- newUndefinedVector n
+    let go (-1) = return v
+        go !k = unsafeWriteVector v k x >> go (k-1 :: Int)
+    go (n-1)
 
 -------------------------------------------------------------------------
 
@@ -153,6 +163,10 @@ readMatrix = safeIndexM unsafeReadMatrix
 writeMatrix :: Storable t => STMatrix s t -> Int -> Int -> t -> ST s ()
 writeMatrix = safeIndexM unsafeWriteMatrix
 
+{-# NOINLINE newUndefinedMatrix #-}
+newUndefinedMatrix :: Element t => MatrixOrder -> Int -> Int -> ST s (STMatrix s t)
+newUndefinedMatrix order r c = unsafeIOToST $ fmap STMatrix $ createMatrix order r c
+
 {-# NOINLINE newMatrix #-}
 newMatrix :: Element t => t -> Int -> Int -> ST s (STMatrix s t)
-newMatrix v r c = unsafeThawMatrix . reshape c . constant v $ r*c
+newMatrix v r c = unsafeThawMatrix $ reshape c $ runSTVector $ newVector v (r*c)
-- 
cgit v1.2.3