module  PreludeArray ( Array, Assoc((:=)), array, listArray, (!), bounds,
 		     indices, elems, assocs, accumArray, (//), accum, amap,
 		     ixmap
 		   ) where
 
 {-#Prelude#-}  -- Indicates definitions of compiler prelude symbols
 
 -- This module uses some simple techniques with updatable vectors to
 -- avoid vector copying in loops where single threading is obvious.
 -- This is rather fragile and depends on the way the compiler handles
 -- strictness.
 
 import PreludeBltinArray
 
 infixl 9  !
 infixl 9  //
 infix  1  :=
 
 data  Assoc a b =  a := b  deriving (Eq, Ord, Ix, Text, Binary)
 data  (Ix a)    => Array a b = MkArray (a,a) {-#STRICT#-}
                                        (Vector (Box b)) {-#STRICT#-}
 				       deriving ()
 
 array		:: (Ix a) => (a,a) -> [Assoc a b] -> Array a b
 listArray	:: (Ix a) => (a,a) -> [b] -> Array a b
 (!)		:: (Ix a) => Array a b -> a -> b
 bounds		:: (Ix a) => Array a b -> (a,a)
 indices		:: (Ix a) => Array a b -> [a]
 elems		:: (Ix a) => Array a b -> [b]
 assocs		:: (Ix a) => Array a b -> [Assoc a b]
 accumArray	:: (Ix a) => (b -> c -> b) -> b -> (a,a) -> [Assoc a c]
 			     -> Array a b
 (//)		:: (Ix a) => Array a b -> [Assoc a b] -> Array a b
 accum		:: (Ix a) => (b -> c -> b) -> Array a b -> [Assoc a c]
 			     -> Array a b
 amap		:: (Ix a) => (b -> c) -> Array a b -> Array a c
 ixmap		:: (Ix a, Ix b) => (a,a) -> (a -> b) -> Array b c
 			     -> Array a c
 
 -- Arrays are a datatype containing a bounds pair and a vector of values.
 -- Uninitialized array elements contain an error value.
 
 -- Primitive vectors now contain only unboxed values.  This permits us to
 -- treat array indexing as an atomic operation without forcing the element
 -- being accessed.  The boxing and unboxing of array elements happens
 -- explicitly using these operations:
 
 data Box a = MkBox a
 unBox (MkBox x) = x
 {-# unBox :: Inline #-}
 
 
 -- Array construction and update using index/value associations share
 -- the same helper function.
 
 array b@(bmin, bmax) ivs =
   let size = (index b bmax) + 1
       v = primMakeVector size uninitializedArrayError
   in (MkArray b (updateArrayIvs b v ivs))
 {-# array :: Inline #-}
 
 a@(MkArray b v) // ivs =
   let v' = primCopyVector v
   in (MkArray b (updateArrayIvs b v' ivs))
 {-# (//) :: Inline #-}
 
 updateArrayIvs b v ivs = 
   let g (i := x) next =  strict1 (primVectorUpdate v (index b i) (MkBox x))
                                  next
   in foldr g v ivs
 {-# updateArrayIvs :: Inline #-}
 
 uninitializedArrayError = 
   MkBox (error "(!){PreludeArray}: uninitialized array element.")
 
 
 -- when mapping a list onto an array, be smart and don't do full index 
 -- computation
 
 listArray b@(bmin, bmax) vs =
   let size = (index b bmax) + 1
       v = primMakeVector size uninitializedArrayError
   in (MkArray b (updateArrayVs size v vs))
 {-# listArray :: Inline #-}
 
 updateArrayVs size v vs =
   let g x next j = if (j == size)
                      then v
 		     else strict1 (primVectorUpdate v j (MkBox x))
 		                  (next (j + 1))
   in foldr g (\ _ -> v) vs 0
 {-# updateArrayVs :: Inline #-}
 
 
 -- Array access
 
 a@(MkArray b v) ! i = unBox (primVectorSel v (index b i))
 {-# (!) :: Inline #-}
 
 bounds (MkArray b _)  = b
 
 indices		      = range . bounds
 
 
 -- Again, when mapping array elements into a list, be smart and don't do 
 -- the full index computation for every element.
 
 elems a@(MkArray b@(bmin, bmax) v) =
   build (\ c n -> 
           let size = (index b bmax) + 1
 	      g j  = if (j == size)
 	                then n
 			else c (unBox (primVectorSel v j)) (g (j + 1))
           -- This strict1 is so size doesn't get inlined and recomputed
 	  -- at every iteration.  It should also force the array argument
 	  -- to be strict.
           in strict1 size (g 0))
 {-# elems :: Inline #-}
 
 assocs a@(MkArray b@(bmin, bmax) v) =
   build (\ c n ->
           let g i next j = let y = unBox (primVectorSel v j)
                            in c (i := y) (next (j + 1))
 	  in foldr g (\ _ -> n) (range b) 0)
 {-# assocs :: Inline #-}
 
 
 -- accum and accumArray share the same helper function.  The difference is
 -- that accum makes a copy of an existing array and accumArray creates
 -- a new one with all elements initialized to the given value.
 
 accum f a@(MkArray b v) ivs =
   let v' = primCopyVector v
   in (MkArray b (accumArrayIvs f b v' ivs))
 {-# accum :: Inline #-}
 
 accumArray f z b@(bmin, bmax) ivs =
   let size = (index b bmax) + 1
       v = primMakeVector size (MkBox z)
   in (MkArray b (accumArrayIvs f b v ivs))
 {-# accumArray :: Inline #-}
 
 
 -- This is a bit tricky.  We need to force the access to the array element
 -- before the update, but not force the thunk that is the value of the
 -- array element unless f is strict.
 
 accumArrayIvs f b v ivs =
   let g (i := x) next = 
         let j = index b i
 	    y = primVectorSel v j
 	in strict1
 	     y
 	     (strict1 (primVectorUpdate v j (MkBox (f (unBox y) x)))
 	              next)
   in foldr g v ivs
 {-# accumArrayIvs :: Inline #-}
 
 
 -- again, be smart and bypass full array indexing on array mapping
 
 amap f a@(MkArray b@(bmin, bmax) v) =
   let size = (index b bmax) + 1
       v' = primMakeVector size uninitializedArrayError
       g j = if (j == size)
               then v'
 	      else let y = primVectorSel v j
 	           in strict1 (primVectorUpdate v' j (MkBox (f (unBox y))))
 	                      (g (j + 1))
   in (MkArray b (g 0))
 {-# amap :: Inline #-}
 
 
 -- can't bypass the index computation here since f needs it as an argument
 
 ixmap b f a           = array b [i := a ! f i | i <- range b]
 {-# ixmap :: Inline #-}
 
 
 -- random other stuff
 
 instance  (Ix a, Eq b)  => Eq (Array a b)  where
     a == a'  	        =  assocs a == assocs a'
 
 instance  (Ix a, Ord b) => Ord (Array a b)  where
     a <=  a'  	    	=  assocs a <=  assocs a'
 
 instance  (Ix a, Text a, Text b) => Text (Array a b)  where
     showsPrec p a = showParen (p > 9) (
 		    showString "array " .
 		    shows (bounds a) . showChar ' ' .
 		    shows (assocs a)                  )
 
     readsPrec p = readParen (p > 9)
 	   (\r -> [(array b as, u) | ("array",s) <- lex r,
 				     (b,t)       <- reads s,
 				     (as,u)      <- reads t   ]
 		  ++
 		  [(listArray b xs, u) | ("listArray",s) <- lex r,
 					 (b,t)           <- reads s,
 					 (xs,u)          <- reads t ])