{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}

module Torch.Index
  ( slice,
    lslice,
  )
where

import Control.Monad ((>=>))
import Data.Void
import Language.Haskell.TH.Lib
import Language.Haskell.TH.Quote (QuasiQuoter (..))
import Language.Haskell.TH.Syntax hiding (Unsafe)
import Text.Megaparsec as M
import Text.Megaparsec.Char hiding (space)
import Text.Megaparsec.Char.Lexer
import Torch.Tensor

type Parser = Parsec Void String

sc :: Parser ()
sc :: Parser ()
sc = Parser () -> Parser () -> Parser () -> Parser ()
forall e s (m :: * -> *).
MonadParsec e s m =>
m () -> m () -> m () -> m ()
space Parser ()
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
m ()
space1 Parser ()
forall a. ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a
empty Parser ()
forall a. ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a
empty

lexm :: Parser a -> Parser a
lexm :: forall a. Parser a -> Parser a
lexm = Parser ()
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall e s (m :: * -> *) a. MonadParsec e s m => m () -> m a -> m a
lexeme Parser ()
sc

parseSlice :: String -> Q [Exp]
parseSlice :: [Char] -> Q [Exp]
parseSlice [Char]
str =
  case Parsec Void [Char] [Exp]
-> [Char] -> [Char] -> Either (ParseErrorBundle [Char] Void) [Exp]
forall e s a.
Parsec e s a -> [Char] -> s -> Either (ParseErrorBundle s e) a
M.runParser Parsec Void [Char] [Exp]
parse' [Char]
"slice" [Char]
str of
    Left ParseErrorBundle [Char] Void
e -> [Char] -> Q [Exp]
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> Q [Exp]) -> [Char] -> Q [Exp]
forall a b. (a -> b) -> a -> b
$ ParseErrorBundle [Char] Void -> [Char]
forall a. Show a => a -> [Char]
show ParseErrorBundle [Char] Void
e
    Right [Exp]
v -> [Exp] -> Q [Exp]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return [Exp]
v
  where
    parse' :: Parser [Exp]
    parse' :: Parsec Void [Char] [Exp]
parse' = (Parser ()
sc Parser ()
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a b.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity b -> ParsecT Void [Char] Identity b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a. Parser a -> Parser a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try ParsecT Void [Char] Identity Exp
slice ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a. Parser a -> Parser a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try ParsecT Void [Char] Identity Exp
bool ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a. Parser a -> Parser a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try ParsecT Void [Char] Identity Exp
other ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ParsecT Void [Char] Identity Exp
number)) ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Char -> Parsec Void [Char] [Exp]
forall (m :: * -> *) a sep. MonadPlus m => m a -> m sep -> m [a]
`sepBy` Token [Char] -> ParsecT Void [Char] Identity (Token [Char])
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
Token s -> m (Token s)
char Char
Token [Char]
','
    other :: Parser Exp
    other :: ParsecT Void [Char] Identity Exp
other =
      ( do
          Tokens [Char]
_ <- Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string ([Char]
Tokens [Char]
"None" :: Tokens String)
          Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE 'None
      )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ( do
                Tokens [Char]
_ <- Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string ([Char]
Tokens [Char]
"Ellipsis" :: Tokens String)
                Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE 'Ellipsis
            )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ( do
                Tokens [Char]
_ <- Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string ([Char]
Tokens [Char]
"..." :: Tokens String)
                Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE 'Ellipsis
            )
    bool :: Parser Exp
    bool :: ParsecT Void [Char] Identity Exp
bool =
      ( do
          Tokens [Char]
_ <- Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
"True"
          Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE 'True
      )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ( do
                Tokens [Char]
_ <- Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
"False"
                Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE 'False
            )
    number :: Parser Exp
    number :: ParsecT Void [Char] Identity Exp
number =
      ( do
          Integer
v <- Parser Integer -> Parser Integer
forall a. Parser a -> Parser a
lexm Parser Integer
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m a
decimal
          Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Lit -> Exp
LitE (Integer -> Lit
IntegerL Integer
v)
      )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ( do
                Tokens [Char]
_ <- Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
"-"
                Integer
v <- Parser Integer -> Parser Integer
forall a. Parser a -> Parser a
lexm Parser Integer
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m a
decimal
                Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Lit -> Exp
LitE (Integer -> Lit
IntegerL (- Integer
v))
            )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ( do
                [Char]
v <- Parser [Char] -> Parser [Char]
forall a. Parser a -> Parser a
lexm (Parser [Char] -> Parser [Char]) -> Parser [Char] -> Parser [Char]
forall a b. (a -> b) -> a -> b
$ ParsecT Void [Char] Identity Char
-> ParsecT Void [Char] Identity Char
-> Parser [Char]
-> Parser [Char]
forall (m :: * -> *) open close a.
Applicative m =>
m open -> m close -> m a -> m a
between (Token [Char] -> ParsecT Void [Char] Identity (Token [Char])
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
Token s -> m (Token s)
char Char
Token [Char]
'{') (Token [Char] -> ParsecT Void [Char] Identity (Token [Char])
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
Token s -> m (Token s)
char Char
Token [Char]
'}') (ParsecT Void [Char] Identity Char -> Parser [Char]
forall (m :: * -> *) a. MonadPlus m => m a -> m [a]
some ParsecT Void [Char] Identity Char
ParsecT Void [Char] Identity (Token [Char])
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
m (Token s)
alphaNumChar)
                Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
VarE ([Char] -> Name
mkName [Char]
v)
            )
    slice :: ParsecT Void [Char] Identity Exp
slice =
      ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a. Parser a -> Parser a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try
        ( do
            Exp
a <- ParsecT Void [Char] Identity Exp
number
            Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
":"
            Exp
b <- ParsecT Void [Char] Identity Exp
number
            Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
":"
            Exp
c <- ParsecT Void [Char] Identity Exp
number
            Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Slice) ([Maybe Exp] -> Exp
TupE [Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
a, Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
b, Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
c])
        )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a. Parser a -> Parser a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try
          ( do
              Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
":"
              Exp
b <- ParsecT Void [Char] Identity Exp
number
              Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
":"
              Exp
c <- ParsecT Void [Char] Identity Exp
number
              Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Slice) ([Maybe Exp] -> Exp
TupE [Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
ConE 'None), Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
b, Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
c])
          )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a. Parser a -> Parser a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try
          ( do
              Exp
a <- ParsecT Void [Char] Identity Exp
number
              Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
"::"
              Exp
c <- ParsecT Void [Char] Identity Exp
number
              Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Slice) ([Maybe Exp] -> Exp
TupE [Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
a, Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
ConE 'None), Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
c])
          )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a. Parser a -> Parser a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try
          ( do
              Exp
a <- ParsecT Void [Char] Identity Exp
number
              Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
":"
              Exp
b <- ParsecT Void [Char] Identity Exp
number
              Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Slice) ([Maybe Exp] -> Exp
TupE [Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
a, Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
b])
          )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a. Parser a -> Parser a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try
          ( do
              Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
"::"
              Exp
c <- ParsecT Void [Char] Identity Exp
number
              Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Slice) ([Maybe Exp] -> Exp
TupE [Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
ConE 'None), Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
ConE 'None), Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
c])
          )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a. Parser a -> Parser a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try
          ( do
              Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
":"
              Exp
b <- ParsecT Void [Char] Identity Exp
number
              Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
":"
              Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Slice) ([Maybe Exp] -> Exp
TupE [Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
ConE 'None), Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
b])
          )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a. Parser a -> Parser a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try
          ( do
              Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
":"
              Exp
b <- ParsecT Void [Char] Identity Exp
number
              Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Slice) ([Maybe Exp] -> Exp
TupE [Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
ConE 'None), Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
b])
          )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a. Parser a -> Parser a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try
          ( do
              Exp
a <- ParsecT Void [Char] Identity Exp
number
              Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
"::"
              Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Slice) ([Maybe Exp] -> Exp
TupE [Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
a, Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
ConE 'None)])
          )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a. Parser a -> Parser a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try
          ( do
              Exp
a <- ParsecT Void [Char] Identity Exp
number
              Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
":"
              Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Slice) ([Maybe Exp] -> Exp
TupE [Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
a, Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
ConE 'None)])
          )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a. Parser a -> Parser a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try
          ( do
              Tokens [Char]
_ <- Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
"::"
              Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Slice) (Name -> Exp
ConE '())
          )
        ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
-> ParsecT Void [Char] Identity Exp
forall a.
ParsecT Void [Char] Identity a
-> ParsecT Void [Char] Identity a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ( do
                Tokens [Char]
_ <- Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a. Parser a -> Parser a
lexm (Parser (Tokens [Char]) -> Parser (Tokens [Char]))
-> Parser (Tokens [Char]) -> Parser (Tokens [Char])
forall a b. (a -> b) -> a -> b
$ Tokens [Char] -> Parser (Tokens [Char])
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
string Tokens [Char]
":"
                Exp -> ParsecT Void [Char] Identity Exp
forall a. a -> ParsecT Void [Char] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ParsecT Void [Char] Identity Exp)
-> Exp -> ParsecT Void [Char] Identity Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Slice) (Name -> Exp
ConE '())
            )

-- | Generate a slice from a [python compatible expression](https://pytorch.org/cppdocs/notes/tensor_indexing.html).
-- When you take the odd-numbered element of tensor with `tensor[1::2]` in python,
-- you can write `tensor ! [slice|1::2|]` in hasktorch.
slice :: QuasiQuoter
slice :: QuasiQuoter
slice =
  QuasiQuoter
    { quoteExp :: [Char] -> Q Exp
quoteExp = [Char] -> Q [Exp]
parseSlice ([Char] -> Q [Exp]) -> ([Exp] -> Q Exp) -> [Char] -> Q Exp
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> [Exp] -> Q Exp
qconcat,
      quotePat :: [Char] -> Q Pat
quotePat = [Char] -> [Char] -> Q Pat
forall a. HasCallStack => [Char] -> a
error [Char]
"quotePat is not implemented for slice.",
      quoteDec :: [Char] -> Q [Dec]
quoteDec = [Char] -> [Char] -> Q [Dec]
forall a. HasCallStack => [Char] -> a
error [Char]
"quoteDec is not implemented for slice.",
      quoteType :: [Char] -> Q Type
quoteType = [Char] -> [Char] -> Q Type
forall a. HasCallStack => [Char] -> a
error [Char]
"quoteType is not implemented for slice."
    }
  where
    qconcat :: [Exp] -> Q Exp
    qconcat :: [Exp] -> Q Exp
qconcat [Exp
exp] = Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
exp
    qconcat [Exp]
exps = Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Maybe Exp] -> Exp
TupE ([Maybe Exp] -> Exp) -> [Maybe Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> Maybe Exp) -> [Exp] -> [Maybe Exp]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Maybe Exp
forall a. a -> Maybe a
Just [Exp]
exps

-- | Generate a lens from a [python compatible expression](https://pytorch.org/cppdocs/notes/tensor_indexing.html).
-- When you take the odd-numbered elements of tensor with `tensor[1::2]` in python,
-- you can write `tensor ^. [lslice|1::2|]` in hasktorch.
-- When you put 2 in the odd numbered elements of the tensor,
-- you can write `tensor & [lslice|1::2|] ~. 2`.
lslice :: QuasiQuoter
lslice :: QuasiQuoter
lslice =
  QuasiQuoter
    { quoteExp :: [Char] -> Q Exp
quoteExp = [Char] -> Q [Exp]
parseSlice ([Char] -> Q [Exp]) -> ([Exp] -> Q Exp) -> [Char] -> Q Exp
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> [Exp] -> Q Exp
qconcat,
      quotePat :: [Char] -> Q Pat
quotePat = [Char] -> [Char] -> Q Pat
forall a. HasCallStack => [Char] -> a
error [Char]
"quotePat is not implemented for slice.",
      quoteDec :: [Char] -> Q [Dec]
quoteDec = [Char] -> [Char] -> Q [Dec]
forall a. HasCallStack => [Char] -> a
error [Char]
"quoteDec is not implemented for slice.",
      quoteType :: [Char] -> Q Type
quoteType = [Char] -> [Char] -> Q Type
forall a. HasCallStack => [Char] -> a
error [Char]
"quoteType is not implemented for slice."
    }
  where
    qconcat :: [Exp] -> Q Exp
    qconcat :: [Exp] -> Q Exp
qconcat [Exp
exp] = Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'toLens) Exp
exp
    qconcat [Exp]
exps = Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'toLens) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ [Maybe Exp] -> Exp
TupE ([Maybe Exp] -> Exp) -> [Maybe Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> Maybe Exp) -> [Exp] -> [Maybe Exp]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Maybe Exp
forall a. a -> Maybe a
Just [Exp]
exps