{- |
Module      : Language.Egison.Type.Index
Licence     : MIT

This module defines tensor index types for the Egison type system.
Indices can be superscript (contravariant, ~i) or subscript (covariant, _i).
-}

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}

module Language.Egison.Type.Index
  ( IndexKind(..)
  , Index(..)
  , IndexSpec
  , IndexTyVar(..)
  , isSupSubPair
  , isSuperscript
  , isSubscript
  , isPlaceholder
  , indexSymbol
  , flipIndexKind
  ) where

import           Data.Hashable (Hashable)
import           GHC.Generics (Generic)

-- | The kind of tensor index
data IndexKind
  = Superscript    -- ^ Contravariant index, written as ~i
  | Subscript      -- ^ Covariant index, written as _i
  deriving (IndexKind -> IndexKind -> Bool
(IndexKind -> IndexKind -> Bool)
-> (IndexKind -> IndexKind -> Bool) -> Eq IndexKind
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: IndexKind -> IndexKind -> Bool
== :: IndexKind -> IndexKind -> Bool
$c/= :: IndexKind -> IndexKind -> Bool
/= :: IndexKind -> IndexKind -> Bool
Eq, Eq IndexKind
Eq IndexKind =>
(IndexKind -> IndexKind -> Ordering)
-> (IndexKind -> IndexKind -> Bool)
-> (IndexKind -> IndexKind -> Bool)
-> (IndexKind -> IndexKind -> Bool)
-> (IndexKind -> IndexKind -> Bool)
-> (IndexKind -> IndexKind -> IndexKind)
-> (IndexKind -> IndexKind -> IndexKind)
-> Ord IndexKind
IndexKind -> IndexKind -> Bool
IndexKind -> IndexKind -> Ordering
IndexKind -> IndexKind -> IndexKind
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: IndexKind -> IndexKind -> Ordering
compare :: IndexKind -> IndexKind -> Ordering
$c< :: IndexKind -> IndexKind -> Bool
< :: IndexKind -> IndexKind -> Bool
$c<= :: IndexKind -> IndexKind -> Bool
<= :: IndexKind -> IndexKind -> Bool
$c> :: IndexKind -> IndexKind -> Bool
> :: IndexKind -> IndexKind -> Bool
$c>= :: IndexKind -> IndexKind -> Bool
>= :: IndexKind -> IndexKind -> Bool
$cmax :: IndexKind -> IndexKind -> IndexKind
max :: IndexKind -> IndexKind -> IndexKind
$cmin :: IndexKind -> IndexKind -> IndexKind
min :: IndexKind -> IndexKind -> IndexKind
Ord, Int -> IndexKind -> ShowS
[IndexKind] -> ShowS
IndexKind -> String
(Int -> IndexKind -> ShowS)
-> (IndexKind -> String)
-> ([IndexKind] -> ShowS)
-> Show IndexKind
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> IndexKind -> ShowS
showsPrec :: Int -> IndexKind -> ShowS
$cshow :: IndexKind -> String
show :: IndexKind -> String
$cshowList :: [IndexKind] -> ShowS
showList :: [IndexKind] -> ShowS
Show, (forall x. IndexKind -> Rep IndexKind x)
-> (forall x. Rep IndexKind x -> IndexKind) -> Generic IndexKind
forall x. Rep IndexKind x -> IndexKind
forall x. IndexKind -> Rep IndexKind x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. IndexKind -> Rep IndexKind x
from :: forall x. IndexKind -> Rep IndexKind x
$cto :: forall x. Rep IndexKind x -> IndexKind
to :: forall x. Rep IndexKind x -> IndexKind
Generic, Eq IndexKind
Eq IndexKind =>
(Int -> IndexKind -> Int)
-> (IndexKind -> Int) -> Hashable IndexKind
Int -> IndexKind -> Int
IndexKind -> Int
forall a. Eq a => (Int -> a -> Int) -> (a -> Int) -> Hashable a
$chashWithSalt :: Int -> IndexKind -> Int
hashWithSalt :: Int -> IndexKind -> Int
$chash :: IndexKind -> Int
hash :: IndexKind -> Int
Hashable)

-- | A tensor index
data Index
  = IndexSym IndexKind String      -- ^ Named index, e.g., _i, ~j
  | IndexPlaceholder IndexKind     -- ^ Placeholder index, e.g., _#, ~#
  | IndexVar String                -- ^ Index variable (for type-level computation)
  deriving (Index -> Index -> Bool
(Index -> Index -> Bool) -> (Index -> Index -> Bool) -> Eq Index
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Index -> Index -> Bool
== :: Index -> Index -> Bool
$c/= :: Index -> Index -> Bool
/= :: Index -> Index -> Bool
Eq, Eq Index
Eq Index =>
(Index -> Index -> Ordering)
-> (Index -> Index -> Bool)
-> (Index -> Index -> Bool)
-> (Index -> Index -> Bool)
-> (Index -> Index -> Bool)
-> (Index -> Index -> Index)
-> (Index -> Index -> Index)
-> Ord Index
Index -> Index -> Bool
Index -> Index -> Ordering
Index -> Index -> Index
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Index -> Index -> Ordering
compare :: Index -> Index -> Ordering
$c< :: Index -> Index -> Bool
< :: Index -> Index -> Bool
$c<= :: Index -> Index -> Bool
<= :: Index -> Index -> Bool
$c> :: Index -> Index -> Bool
> :: Index -> Index -> Bool
$c>= :: Index -> Index -> Bool
>= :: Index -> Index -> Bool
$cmax :: Index -> Index -> Index
max :: Index -> Index -> Index
$cmin :: Index -> Index -> Index
min :: Index -> Index -> Index
Ord, Int -> Index -> ShowS
[Index] -> ShowS
Index -> String
(Int -> Index -> ShowS)
-> (Index -> String) -> ([Index] -> ShowS) -> Show Index
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Index -> ShowS
showsPrec :: Int -> Index -> ShowS
$cshow :: Index -> String
show :: Index -> String
$cshowList :: [Index] -> ShowS
showList :: [Index] -> ShowS
Show, (forall x. Index -> Rep Index x)
-> (forall x. Rep Index x -> Index) -> Generic Index
forall x. Rep Index x -> Index
forall x. Index -> Rep Index x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Index -> Rep Index x
from :: forall x. Index -> Rep Index x
$cto :: forall x. Rep Index x -> Index
to :: forall x. Rep Index x -> Index
Generic, Eq Index
Eq Index =>
(Int -> Index -> Int) -> (Index -> Int) -> Hashable Index
Int -> Index -> Int
Index -> Int
forall a. Eq a => (Int -> a -> Int) -> (a -> Int) -> Hashable a
$chashWithSalt :: Int -> Index -> Int
hashWithSalt :: Int -> Index -> Int
$chash :: Index -> Int
hash :: Index -> Int
Hashable)

-- | A sequence of indices
type IndexSpec = [Index]

-- | Index variable for type schemes (type-level)
newtype IndexTyVar = IndexTyVar String
  deriving (IndexTyVar -> IndexTyVar -> Bool
(IndexTyVar -> IndexTyVar -> Bool)
-> (IndexTyVar -> IndexTyVar -> Bool) -> Eq IndexTyVar
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: IndexTyVar -> IndexTyVar -> Bool
== :: IndexTyVar -> IndexTyVar -> Bool
$c/= :: IndexTyVar -> IndexTyVar -> Bool
/= :: IndexTyVar -> IndexTyVar -> Bool
Eq, Eq IndexTyVar
Eq IndexTyVar =>
(IndexTyVar -> IndexTyVar -> Ordering)
-> (IndexTyVar -> IndexTyVar -> Bool)
-> (IndexTyVar -> IndexTyVar -> Bool)
-> (IndexTyVar -> IndexTyVar -> Bool)
-> (IndexTyVar -> IndexTyVar -> Bool)
-> (IndexTyVar -> IndexTyVar -> IndexTyVar)
-> (IndexTyVar -> IndexTyVar -> IndexTyVar)
-> Ord IndexTyVar
IndexTyVar -> IndexTyVar -> Bool
IndexTyVar -> IndexTyVar -> Ordering
IndexTyVar -> IndexTyVar -> IndexTyVar
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: IndexTyVar -> IndexTyVar -> Ordering
compare :: IndexTyVar -> IndexTyVar -> Ordering
$c< :: IndexTyVar -> IndexTyVar -> Bool
< :: IndexTyVar -> IndexTyVar -> Bool
$c<= :: IndexTyVar -> IndexTyVar -> Bool
<= :: IndexTyVar -> IndexTyVar -> Bool
$c> :: IndexTyVar -> IndexTyVar -> Bool
> :: IndexTyVar -> IndexTyVar -> Bool
$c>= :: IndexTyVar -> IndexTyVar -> Bool
>= :: IndexTyVar -> IndexTyVar -> Bool
$cmax :: IndexTyVar -> IndexTyVar -> IndexTyVar
max :: IndexTyVar -> IndexTyVar -> IndexTyVar
$cmin :: IndexTyVar -> IndexTyVar -> IndexTyVar
min :: IndexTyVar -> IndexTyVar -> IndexTyVar
Ord, Int -> IndexTyVar -> ShowS
[IndexTyVar] -> ShowS
IndexTyVar -> String
(Int -> IndexTyVar -> ShowS)
-> (IndexTyVar -> String)
-> ([IndexTyVar] -> ShowS)
-> Show IndexTyVar
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> IndexTyVar -> ShowS
showsPrec :: Int -> IndexTyVar -> ShowS
$cshow :: IndexTyVar -> String
show :: IndexTyVar -> String
$cshowList :: [IndexTyVar] -> ShowS
showList :: [IndexTyVar] -> ShowS
Show, (forall x. IndexTyVar -> Rep IndexTyVar x)
-> (forall x. Rep IndexTyVar x -> IndexTyVar) -> Generic IndexTyVar
forall x. Rep IndexTyVar x -> IndexTyVar
forall x. IndexTyVar -> Rep IndexTyVar x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. IndexTyVar -> Rep IndexTyVar x
from :: forall x. IndexTyVar -> Rep IndexTyVar x
$cto :: forall x. Rep IndexTyVar x -> IndexTyVar
to :: forall x. Rep IndexTyVar x -> IndexTyVar
Generic)

-- | Check if two indices form a superscript-subscript pair (for contraction)
-- For example, ~i and _i form a pair
isSupSubPair :: Index -> Index -> Bool
isSupSubPair :: Index -> Index -> Bool
isSupSubPair (IndexSym IndexKind
Superscript String
s1) (IndexSym IndexKind
Subscript String
s2) = String
s1 String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
s2
isSupSubPair (IndexSym IndexKind
Subscript String
s1) (IndexSym IndexKind
Superscript String
s2) = String
s1 String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
s2
isSupSubPair Index
_ Index
_ = Bool
False

-- | Check if an index is a superscript
isSuperscript :: Index -> Bool
isSuperscript :: Index -> Bool
isSuperscript (IndexSym IndexKind
Superscript String
_)   = Bool
True
isSuperscript (IndexPlaceholder IndexKind
Superscript) = Bool
True
isSuperscript Index
_                          = Bool
False

-- | Check if an index is a subscript
isSubscript :: Index -> Bool
isSubscript :: Index -> Bool
isSubscript (IndexSym IndexKind
Subscript String
_)   = Bool
True
isSubscript (IndexPlaceholder IndexKind
Subscript) = Bool
True
isSubscript Index
_                        = Bool
False

-- | Check if an index is a placeholder
isPlaceholder :: Index -> Bool
isPlaceholder :: Index -> Bool
isPlaceholder (IndexPlaceholder IndexKind
_) = Bool
True
isPlaceholder Index
_                    = Bool
False

-- | Get the symbol name from an index (if it has one)
indexSymbol :: Index -> Maybe String
indexSymbol :: Index -> Maybe String
indexSymbol (IndexSym IndexKind
_ String
s) = String -> Maybe String
forall a. a -> Maybe a
Just String
s
indexSymbol (IndexVar String
s)   = String -> Maybe String
forall a. a -> Maybe a
Just String
s
indexSymbol Index
_              = Maybe String
forall a. Maybe a
Nothing

-- | Flip the kind of an index (superscript <-> subscript)
flipIndexKind :: IndexKind -> IndexKind
flipIndexKind :: IndexKind -> IndexKind
flipIndexKind IndexKind
Superscript = IndexKind
Subscript
flipIndexKind IndexKind
Subscript   = IndexKind
Superscript