From 8ebb805a1ca0cea64a5e0e6e538573690aab2c4d Mon Sep 17 00:00:00 2001
From: Utku Demir <me@utdemir.com>
Date: Fri, 20 Aug 2021 01:25:59 +1200
Subject: [PATCH] Introduce Data.VM.Linear and
 Data.LArray.Mutable.Unlifted.Linear

---
 linear-base.cabal                             |   5 +
 src/Data/LArray/Mutable/Unlifted/Linear.hs    |  29 +++
 .../Mutable/Unlifted/Linear/Internal.hs       | 194 ++++++++++++++++++
 src/Data/V/Linear/Internal/V.hs               |   5 +
 src/Data/VM/Linear.hs                         |  14 ++
 src/Data/VM/Linear/Internal/Instances.hs      |  45 ++++
 src/Data/VM/Linear/Internal/VM.hs             |  90 ++++++++
 7 files changed, 382 insertions(+)
 create mode 100644 src/Data/LArray/Mutable/Unlifted/Linear.hs
 create mode 100644 src/Data/LArray/Mutable/Unlifted/Linear/Internal.hs
 create mode 100644 src/Data/VM/Linear.hs
 create mode 100644 src/Data/VM/Linear/Internal/Instances.hs
 create mode 100644 src/Data/VM/Linear/Internal/VM.hs

diff --git a/linear-base.cabal b/linear-base.cabal
index 3771d360..22a268ce 100644
--- a/linear-base.cabal
+++ b/linear-base.cabal
@@ -39,6 +39,8 @@ library
     Data.Array.Mutable.Linear
     Data.Array.Mutable.Linear.Internal
     Data.Array.Mutable.Unlifted.Linear
+    Data.LArray.Mutable.Unlifted.Linear
+    Data.LArray.Mutable.Unlifted.Linear.Internal
     Data.Array.Polarized
     Data.Array.Polarized.Pull
     Data.Array.Polarized.Pull.Internal
@@ -78,6 +80,9 @@ library
     Data.V.Linear
     Data.V.Linear.Internal.V
     Data.V.Linear.Internal.Instances
+    Data.VM.Linear
+    Data.VM.Linear.Internal.VM
+    Data.VM.Linear.Internal.Instances
     Data.Vector.Mutable.Linear
     Data.Vector.Mutable.Linear.Internal
     Debug.Trace.Linear
diff --git a/src/Data/LArray/Mutable/Unlifted/Linear.hs b/src/Data/LArray/Mutable/Unlifted/Linear.hs
new file mode 100644
index 00000000..8fcda37e
--- /dev/null
+++ b/src/Data/LArray/Mutable/Unlifted/Linear.hs
@@ -0,0 +1,29 @@
+{-# LANGUAGE NoImplicitPrelude #-}
+{-# LANGUAGE MagicHash #-}
+
+-- |
+-- This module provides an unlifted mutable array with a pure
+-- interface. The array elements are stored linearly, so operations
+-- usually require the either a linear callback or a 'Consumable' or
+-- 'Dupable' constraint.
+--
+-- Accessing out-of-bounds indices causes undefined behaviour.
+--
+-- This module is meant to be imported qualified.
+module Data.LArray.Mutable.Unlifted.Linear
+  ( LArray#
+  , alloc
+  , fromList
+  , allocBeside
+  , lseq
+  , size
+  , get
+  , set
+  , update
+  , map
+  , toList
+  , append
+  , dup2
+  ) where
+
+import Data.LArray.Mutable.Unlifted.Linear.Internal
diff --git a/src/Data/LArray/Mutable/Unlifted/Linear/Internal.hs b/src/Data/LArray/Mutable/Unlifted/Linear/Internal.hs
new file mode 100644
index 00000000..6d90f3b9
--- /dev/null
+++ b/src/Data/LArray/Mutable/Unlifted/Linear/Internal.hs
@@ -0,0 +1,194 @@
+{-# OPTIONS_HADDOCK hide #-}
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE LinearTypes #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE UnliftedNewtypes #-}
+
+module Data.LArray.Mutable.Unlifted.Linear.Internal where
+
+import Data.Unrestricted.Linear hiding (lseq, dup2)
+import qualified Data.Unrestricted.Linear
+import Prelude (Int)
+import Data.List (length)
+import qualified Prelude as Prelude
+import Prelude.Linear.Internal
+import qualified Unsafe.Linear as Unsafe
+import qualified GHC.Exts as GHC
+
+-- | A mutable array holding @a@s
+newtype LArray# a = LArray# (GHC.MutableArray# GHC.RealWorld a)
+
+-- | Consume a 'LArray#' and its elements. O(n)
+--
+-- Note that we can not implement a 'Consumable' instance because 'LArray#'
+-- is unlifted.
+lseq :: Consumable a => LArray# a %1-> b %1-> b
+lseq arr b =
+  (Unsafe.toLinear2 (\x _ -> x) b)
+    (map consume arr)
+
+-- | Allocate a mutable array of given size using a default value.
+--
+-- The size should be non-negative.
+alloc :: Int -> a -> (LArray# a %1-> Ur b) %1-> Ur b
+alloc (GHC.I# s) a f =
+  let new = GHC.runRW# Prelude.$ \st ->
+        case GHC.newArray# s a st of
+          (# _, arr #) -> LArray# arr
+   in f new
+{-# NOINLINE alloc #-}  -- prevents runRW# from floating outwards
+
+fromList :: [a] %1-> (LArray# a %1-> Ur b) %1-> Ur b
+fromList = Unsafe.toLinear2 $ \l f ->
+  let !(GHC.I# len) = length l
+      new = GHC.runRW# Prelude.$ \st ->
+        case GHC.newArray# len Prelude.undefined st of
+          (# st', arr #) -> go 0# arr st' l
+   in f new
+ where
+  go :: GHC.Int# -> GHC.MutableArray# GHC.RealWorld a -> GHC.State# GHC.RealWorld -> [a] -> LArray# a
+  go _ arr !_ [] = LArray# arr
+  go i arr !st (x:xs) =
+    let st' = GHC.writeArray# arr i x st
+    in  go (i GHC.+# 1#) arr st' xs
+
+{-# NOINLINE fromList #-}  -- prevents runRW# from floating outwards
+
+-- For the reasoning behind these NOINLINE pragmas, see the discussion at:
+-- https://github.com/tweag/linear-base/pull/187#pullrequestreview-489183531
+
+-- | Allocate a mutable array of given size using a default value,
+-- using another 'LArray#' as a uniqueness proof.
+--
+-- The size should be non-negative.
+allocBeside :: Int -> a -> LArray# b %1-> (# LArray# a, LArray# b #)
+allocBeside (GHC.I# s) a orig =
+  let new = GHC.runRW# Prelude.$ \st ->
+        case GHC.newArray# s a st of
+          (# _, arr #) -> LArray# arr
+   in (# new, orig #)
+{-# NOINLINE allocBeside #-}  -- prevents runRW# from floating outwards
+
+size :: LArray# a %1-> (# Ur Int, LArray# a #)
+size = Unsafe.toLinear go
+  where
+    go :: LArray# a -> (# Ur Int, LArray# a #)
+    go (LArray# arr) =
+      let !s = GHC.sizeofMutableArray# arr
+      in  (# Ur (GHC.I# s), LArray# arr  #)
+
+get :: Dupable a => Int -> LArray# a %1-> (# a, LArray# a #)
+get ix arr = runRW# (update' ix Data.Unrestricted.Linear.dup2 arr)
+{-# NOINLINE get #-}  -- prevents the runRW# effect from being reordered
+
+set :: Consumable a => Int -> a %1-> LArray# a %1-> LArray# a
+set ix x xs =
+  (\(# (), r #) -> r)
+    (runRW# (update' ix (\old -> (consume old, x)) xs))
+{-# NOINLINE set #-}  -- prevents the runRW# effect from being reordered
+
+update' :: Int -> (a %1-> (b, a)) %1-> LArray# a %1-> GHC.State# GHC.RealWorld %1-> (# b, LArray# a #)
+update' (GHC.I# ix) = Unsafe.toLinear3 go
+  where
+    go :: (a %1-> (b, a)) -> LArray# a -> GHC.State# GHC.RealWorld -> (# b, LArray# a #)
+    go f (LArray# arr) st =
+      case GHC.readArray# arr ix st of
+        (# st', a #) ->
+          case f a of
+            (b, a') ->
+               case GHC.writeArray# arr ix a' st' of
+                 !_ -> (# b, LArray# arr #)
+{-# INLINE update' #-}
+
+update :: Int -> (a %1-> (b, a)) %1-> LArray# a %1-> (# b, LArray# a #)
+update ix f arr = runRW# (update' ix f arr)
+{-# NOINLINE update #-}  -- prevents the runRW# effect from being reordered
+
+-- | Map over the LArray in-place.
+map :: (a %1-> b) -> LArray# a %1-> LArray# b
+map (f :: a %1-> b) = Unsafe.toLinear (\(LArray# as) ->
+  let -- We alias the input array to write the resulting -- 'b's to,
+      -- just to make the typechecker happy. Care must be taken to
+      -- only read indices from 'as' that is not yet written to 'bs'.
+      bs :: GHC.MutableArray# GHC.RealWorld b
+      bs = GHC.unsafeCoerce# as
+      len :: GHC.Int#
+      len = GHC.sizeofMutableArray# as
+
+      -- For each index ([0..len]), we read the element on 'as', pass
+      -- it through 'f' and write to the same location on 'bs'.
+      go :: GHC.Int# -> GHC.State# GHC.RealWorld -> ()
+      go i st
+        | GHC.I# i Prelude.== GHC.I# len = ()
+        | Prelude.otherwise =
+          case GHC.readArray# as i st of
+            (# st', a #) ->
+              case GHC.writeArray# bs i (f a) st' of
+                !st'' -> go (i GHC.+# 1#) st''
+   in GHC.runRW# (go 0#) `GHC.seq` LArray# bs
+  )
+{-# NOINLINE map #-}
+
+append :: LArray# a %1-> LArray# a %1-> LArray# a
+append (LArray# left) (LArray# right) = Unsafe.toLinear2 go left right
+ where
+  go l r =
+    let lsize = GHC.sizeofMutableArray# l
+        rsize = GHC.sizeofMutableArray# r
+     in GHC.runRW# (\st ->
+          case GHC.newArray# (lsize GHC.+# rsize) Prelude.undefined st of
+            (# st', dst #) ->
+              case GHC.copyMutableArray# l 0# dst 0# lsize st' of
+                !st'' ->
+                  case GHC.copyMutableArray# r 0# dst lsize rsize st'' of
+                    !_ -> LArray# dst
+        )
+
+-- | Return the array elements as a lazy list.
+toList :: LArray# a %1-> [a]
+toList (LArray# arr) =
+  Unsafe.toLinear
+    (\xs -> go 0 (GHC.I# (GHC.sizeofMutableArray# xs)) xs)
+    arr
+ where
+  go :: Int -> Int -> GHC.MutableArray# GHC.RealWorld a -> [a]
+  go i len xs
+    | i Prelude.== len = []
+    | GHC.I# i# <- i =
+        case GHC.runRW# (GHC.readArray# xs i#) of
+          (# _, x #) -> x : go (i Prelude.+ 1) len xs
+
+-- | Clone an array and all its elements.
+dup2 :: forall a. Dupable a => LArray# a %1-> (# LArray# a, LArray# a #)
+dup2 = Unsafe.toLinear (\arr -> GHC.runRW# (go arr))
+ where
+  go :: LArray# a -> GHC.State# GHC.RealWorld -> (# LArray# a, LArray# a #)
+  go (LArray# orig) st =
+    let len :: GHC.Int#
+        len = GHC.sizeofMutableArray# orig
+     in case GHC.newArray# len Prelude.undefined st of
+          (# st', left #) ->
+            case GHC.newArray# len Prelude.undefined st' of
+              (# st'', right #) ->
+                 let loop :: GHC.Int# -> GHC.State# GHC.RealWorld -> ()
+                     loop i !st0
+                       | GHC.I# i Prelude.== GHC.I# len = ()
+                       | Prelude.otherwise =
+                           case GHC.readArray# orig i st0 of
+                             (# st1, a #) -> case Data.Unrestricted.Linear.dup2 a of
+                               (a1, a2) -> case GHC.writeArray# left i a1 st1 of
+                                 st2 -> case GHC.writeArray# right i a2 st2 of
+                                   st3 -> loop (i GHC.+# 1#) st3
+                  in case loop 0# st'' of
+                      () -> (# LArray# left, LArray# right #)
+{-# NOINLINE dup2 #-}
+
+-- Utils
+
+runRW# :: forall (r :: GHC.RuntimeRep) (o :: GHC.TYPE r). (GHC.State# GHC.RealWorld %1-> o) %1-> o
+runRW# = Unsafe.coerce GHC.runRW#
diff --git a/src/Data/V/Linear/Internal/V.hs b/src/Data/V/Linear/Internal/V.hs
index 190e9470..980eb4ea 100644
--- a/src/Data/V/Linear/Internal/V.hs
+++ b/src/Data/V/Linear/Internal/V.hs
@@ -23,6 +23,11 @@ module Data.V.Linear.Internal.V
   , iterate
   -- * Type-level utilities
   , caseNat
+  , predNat
+  , expandFunN
+  , contractFunN
+  , continue
+  , Dict (..)
   ) where
 
 import Data.Kind (Type)
diff --git a/src/Data/VM/Linear.hs b/src/Data/VM/Linear.hs
new file mode 100644
index 00000000..95ef3534
--- /dev/null
+++ b/src/Data/VM/Linear.hs
@@ -0,0 +1,14 @@
+{-# OPTIONS_GHC -Wno-dodgy-exports #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+
+module Data.VM.Linear
+  ( VM
+  , elim
+  , make
+  , update
+  , module Data.VM.Linear.Internal.Instances
+  ) where
+
+import Data.VM.Linear.Internal.VM
+import Data.VM.Linear.Internal.Instances ()
+
diff --git a/src/Data/VM/Linear/Internal/Instances.hs b/src/Data/VM/Linear/Internal/Instances.hs
new file mode 100644
index 00000000..2448f5ea
--- /dev/null
+++ b/src/Data/VM/Linear/Internal/Instances.hs
@@ -0,0 +1,45 @@
+{-# OPTIONS_HADDOCK hide #-}
+{-# OPTIONS -Wno-orphans #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE UnboxedTuples #-}
+
+-- | This module contains all instances for VM
+--
+module Data.VM.Linear.Internal.Instances where
+
+import Data.VM.Linear.Internal.VM
+import qualified Data.Functor.Linear.Internal.Functor as Data
+import qualified Data.LArray.Mutable.Unlifted.Linear as LArray
+import Data.Unrestricted.Linear
+
+
+-- # Instances of VM
+-------------------------------------------------------------------------------
+
+instance Data.Functor (VM n) where
+  fmap f (VM xs) = VM (LArray.map f xs)
+
+-- TODO: This requires an efficient 'zip', and which in turns requires
+-- something similar to our Pull Arrays, but with linear elements.
+--
+--   instance Data.Applicative (V n)
+
+-- TODO: This should be possible, but I could not find a way to implement
+-- 'LArray.traverse'.
+-- instance Data.Traversable (VM n)
+--   traverse f (VM arr) =
+--     (\(LArray.LArray arr') -> VM arr')
+--       Data.<$> LArray.traverse f arr
+
+instance Consumable a => Consumable (VM n a) where
+  consume (VM xs) = xs `LArray.lseq` ()
+
+-- TODO: Decide whether we should have Dupable in terms of `VM` or `V`. This
+-- will require moving things around.
+instance Dupable a => Dupable (VM n a) where
+  -- TODO: There must be a better way to implement this using dupN.
+  dup2 (VM xs) =
+    (\(# x, y #) -> (VM x, VM y))
+      (LArray.dup2 xs)
+
diff --git a/src/Data/VM/Linear/Internal/VM.hs b/src/Data/VM/Linear/Internal/VM.hs
new file mode 100644
index 00000000..581e82a2
--- /dev/null
+++ b/src/Data/VM/Linear/Internal/VM.hs
@@ -0,0 +1,90 @@
+{-# OPTIONS_HADDOCK hide #-}
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE LinearTypes #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UnboxedTuples #-}
+module Data.VM.Linear.Internal.VM where
+
+import Data.Kind (Type)
+import Data.Type.Equality
+import GHC.TypeLits
+import Prelude.Linear.Internal
+import Prelude (Either(..) , error)
+import Data.LArray.Mutable.Unlifted.Linear (LArray#)
+import qualified Data.LArray.Mutable.Unlifted.Linear as LArray
+import Data.Unrestricted.Internal.Ur
+import Data.V.Linear.Internal.V (theLength, caseNat, predNat, FunN, expandFunN,
+                                 contractFunN, continue, Dict (..))
+
+{- Developers Note
+
+See the "Developers Note" in Data.V.Linear for an explanation of this module
+structure.
+-}
+
+-- # Type Definitions
+-------------------------------------------------------------------------------
+
+data VM (n :: Nat) (a :: Type) = VM (LArray# a)
+
+-- # API
+-------------------------------------------------------------------------------
+
+consumeV :: VM 0 a %1-> b %1-> b
+consumeV (VM arr) b =
+  LArray.map (error "impossible: arr should be empty" :: a %1-> ()) arr
+    `LArray.lseq` b
+
+{-
+split :: 1 <= n => V n a %1-> (# a, V (n-1) a #)
+split = Unsafe.toLinear split'
+  where
+    split' :: 1 <= n => V n a -> (# a, V (n-1) a #)
+    split' (V xs) = (# Vector.head xs, V (Vector.tail xs) #)
+-}
+
+make :: forall n a b. KnownNat n => FunN n a ((VM n a %1-> Ur b) %1-> Ur b)
+make =
+  continue
+    @n @a @[a] @((VM n a %1-> Ur b) %1-> Ur b)
+    (\xs f -> LArray.fromList xs (\arr -> f (VM arr)))
+    (go @n @a)
+ where
+  go :: forall n' a'. KnownNat n' => FunN n' a' [a']
+  go =
+    case caseNat @n' of
+      Left Refl -> []
+      Right Refl ->
+        contractFunN @n' @a' @[a'] (\a ->
+          case predNat @n' of
+            Dict -> continue @(n'-1) @a' @([a']) (a:) (go @(n'-1) @a'))
+
+update :: forall ix n a b. (KnownNat ix, ix <= n - 1) => (a %1-> (b, a)) %1-> VM n a %1-> (b, VM n a)
+update f (VM arr) =
+  (\(# b, arr' #) -> (b, VM arr'))
+    (LArray.update (theLength @ix) f arr)
+
+elim :: forall n a b. KnownNat n => VM n a %1-> FunN n a b %1-> b
+elim (VM arr) f = go @n @a @b (LArray.toList arr) f
+ where
+  -- invariant: length xs == n
+  go :: forall n' a' b'. KnownNat n' => [a'] %1-> FunN n' a' b' %1-> b'
+  go xs' g =
+    (caseNat @n', xs') & \case
+      (Left Refl, []) -> g
+      (Right Refl, x:xs) ->
+        predNat @n' & \Dict ->
+          go @(n'-1) xs (expandFunN @n' @a' @b' g x)
+      other -> error "invariant violation" other g
+
+toList :: VM n a %1-> [a]
+toList (VM arr) = LArray.toList arr