existensial-slice-hack.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

import Data.Distributive
import Data.Functor.Rep
import Data.List ((!!))
import Data.Singletons
import Data.Singletons.Prelude
import Data.Singletons.TH (promote)
import Data.Singletons.Prelude.List
import Data.Singletons.TypeLits
import GHC.Exts
import GHC.Show
import GHC.TypeLits
import Protolude hiding (show, (<.>))
import qualified Data.Vector as V
import           Data.Foldable
import           Data.Kind
import           Data.Maybe

newtype Tensor (r::[Nat]) a = Tensor { v :: V.Vector a }
    deriving (Functor, Eq, Foldable)

data ProdMap :: (a -> b -> Type) -> [a] -> [b] -> Type where
    PMZ :: ProdMap f '[] '[]
    PMS :: f a b -> ProdMap f as bs -> ProdMap f (a ': as) (b ': bs)

data Slice :: Nat -> Nat -> Type where
    Slice :: Sing l -> Sing c -> Sing r -> Slice (l + c + r) c

slice
    :: (SingI ns, SingI ms)
    => ProdMap Slice ns ms
    -> Tensor ns a
    -> Tensor ms a
slice = undefined

-- given a type-level list `ns` of the number of items from each dimension to take,
-- returns the `ProdMap Slice ms ns` that encodes that.
sliceHeads :: forall ms ns. Sing ns -> ProdMap Slice (ms::[Nat]) (ns::[Nat])
sliceHeads = \case
    SNil -> PMZ
    s@SNat `SCons` ss -> Slice (SNat @0) s meh `PMS` sliceHeads ss
    -- meh has to be :: Sing (m - n), and positive, so in real life, we'd
    -- have to also iterate over ms as well.
      where
        meh = undefined

headsFromList
    :: SingI ms
    => [Integer]
    -> Tensor ms a
    -> (forall ns. SingI ns => Tensor ns a -> r)
    -> r
headsFromList ns t f = withSomeSing ns $ \nsSing ->
                       withSingI nsSing $
                       f (slice (sliceHeads nsSing) t)