{-# LANGUAGE GADTs #-}

module Unison.Runtime.Vector where

import Data.MemoCombinators qualified as Memo
import Data.Vector.Unboxed qualified as UV
import Unison.Prelude

-- A `Vec a` denotes a `Nat -> Maybe a`
data Vec a where
  Scalar :: a -> Vec a
  Vec :: (UV.Unbox a) => UV.Vector a -> Vec a
  Pair :: Vec a -> Vec b -> Vec (a, b)
  Choose :: Vec Bool -> Vec a -> Vec a -> Vec a
  Mux :: Vec Nat -> Vec (Vec a) -> Vec a

-- todo: maybe make representation `(UV.Vector Nat -> UnboxedMap Nat a, Bound)`
-- `UnboxedMap Nat a = (UV.Vector Nat, UV.Vector a)`
-- UnboxedMap Nat could be implemented as an `UArray`
-- `Bound` is Nat, max possible index
-- then easy to implement `+`, `-`, etc

type Nat = Word64

mu :: Vec a -> Nat -> Maybe a
mu :: forall a. Vec a -> Nat -> Maybe a
mu Vec a
v = case Vec a
v of
  Scalar a
a -> Maybe a -> Nat -> Maybe a
forall a b. a -> b -> a
const (a -> Maybe a
forall a. a -> Maybe a
Just a
a)
  Vec Vector a
vs -> \Nat
i -> Vector a
vs Vector a -> Int -> Maybe a
forall a. Unbox a => Vector a -> Int -> Maybe a
UV.!? Nat -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Nat
i
  Choose Vec Bool
cond Vec a
t Vec a
f ->
    let (Nat -> Maybe Bool
condr, Nat -> Maybe a
tr, Nat -> Maybe a
tf) = (Vec Bool -> Nat -> Maybe Bool
forall a. Vec a -> Nat -> Maybe a
mu Vec Bool
cond, Vec a -> Nat -> Maybe a
forall a. Vec a -> Nat -> Maybe a
mu Vec a
t, Vec a -> Nat -> Maybe a
forall a. Vec a -> Nat -> Maybe a
mu Vec a
f)
     in \Nat
i -> Nat -> Maybe Bool
condr Nat
i Maybe Bool -> (Bool -> Maybe a) -> Maybe a
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Bool
b -> if Bool
b then Nat -> Maybe a
tr Nat
i else Nat -> Maybe a
tf Nat
i
  Mux Vec Nat
mux Vec (Vec a)
branches ->
    let muxr :: Nat -> Maybe Nat
muxr = Vec Nat -> Nat -> Maybe Nat
forall a. Vec a -> Nat -> Maybe a
mu Vec Nat
mux
        branchesr :: Nat -> Maybe (Nat -> Maybe a)
branchesr = (Nat -> Maybe (Nat -> Maybe a)) -> Nat -> Maybe (Nat -> Maybe a)
forall a. Integral a => Memo a
Memo Nat
Memo.integral ((Nat -> Maybe (Nat -> Maybe a)) -> Nat -> Maybe (Nat -> Maybe a))
-> (Nat -> Maybe (Nat -> Maybe a)) -> Nat -> Maybe (Nat -> Maybe a)
forall a b. (a -> b) -> a -> b
$ let f :: Nat -> Maybe (Vec a)
f = Vec (Vec a) -> Nat -> Maybe (Vec a)
forall a. Vec a -> Nat -> Maybe a
mu Vec (Vec a)
branches in \Nat
i -> Vec a -> Nat -> Maybe a
forall a. Vec a -> Nat -> Maybe a
mu (Vec a -> Nat -> Maybe a)
-> Maybe (Vec a) -> Maybe (Nat -> Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Nat -> Maybe (Vec a)
f Nat
i
     in \Nat
i -> do Nat
j <- Nat -> Maybe Nat
muxr Nat
i; Nat -> Maybe a
b <- Nat -> Maybe (Nat -> Maybe a)
branchesr Nat
j; Nat -> Maybe a
b Nat
i
  Pair Vec a
v1 Vec b
v2 ->
    let (Nat -> Maybe a
v1r, Nat -> Maybe b
v2r) = (Vec a -> Nat -> Maybe a
forall a. Vec a -> Nat -> Maybe a
mu Vec a
v1, Vec b -> Nat -> Maybe b
forall a. Vec a -> Nat -> Maybe a
mu Vec b
v2)
     in \Nat
i -> (a -> b -> a) -> Maybe a -> Maybe b -> Maybe a
forall a b c. (a -> b -> c) -> Maybe a -> Maybe b -> Maybe c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (,) (Nat -> Maybe a
v1r Nat
i) (Nat -> Maybe b
v2r Nat
i)

-- Returns the maximum `Nat` for which `mu v` may return `Just`.
bound :: Nat -> Vec a -> Nat
bound :: forall a. Nat -> Vec a -> Nat
bound Nat
width Vec a
v = case Vec a
v of
  Scalar a
_ -> Nat
width
  Vec Vector a
vs -> Int -> Nat
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Nat) -> Int -> Nat
forall a b. (a -> b) -> a -> b
$ Vector a -> Int
forall a. Unbox a => Vector a -> Int
UV.length Vector a
vs
  Pair Vec a
v1 Vec b
v2 -> Nat -> Vec a -> Nat
forall a. Nat -> Vec a -> Nat
bound Nat
width Vec a
v1 Nat -> Nat -> Nat
forall a. Ord a => a -> a -> a
`min` Nat -> Vec b -> Nat
forall a. Nat -> Vec a -> Nat
bound Nat
width Vec b
v2
  Choose Vec Bool
cond Vec a
_ Vec a
_ -> Nat -> Vec Bool -> Nat
forall a. Nat -> Vec a -> Nat
bound Nat
width Vec Bool
cond
  Mux Vec Nat
mux Vec (Vec a)
_ -> Nat -> Vec Nat -> Nat
forall a. Nat -> Vec a -> Nat
bound Nat
width Vec Nat
mux

toList :: Vec a -> [a]
toList :: forall a. Vec a -> [a]
toList Vec a
v =
  let n :: Nat
n = Nat -> Vec a -> Nat
forall a. Nat -> Vec a -> Nat
bound Nat
forall a. Bounded a => a
maxBound Vec a
v
      muv :: Nat -> Maybe a
muv = Vec a -> Nat -> Maybe a
forall a. Vec a -> Nat -> Maybe a
mu Vec a
v
   in [Maybe a] -> [a]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe a] -> [a]) -> [Maybe a] -> [a]
forall a b. (a -> b) -> a -> b
$ Nat -> Maybe a
muv (Nat -> Maybe a) -> [Nat] -> [Maybe a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Nat
0 .. Nat
n]