<commit_before>{-# LANGUAGE FlexibleInstances  #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell    #-}

module Database.DSH.FKL.Lang where

import           Prelude hiding ((<$>))
import           Text.PrettyPrint.ANSI.Leijen
import           Text.Printf

import           Database.DSH.Common.Impossible
import           Database.DSH.Common.Pretty
import           Database.DSH.Common.Nat
import qualified Database.DSH.Common.Lang   as L
import           Database.DSH.Common.Type   (Type, Typed, typeOf)


-- | 'LiftedN' defines an FKL dialect in which primitives and
-- operators might be lifted to arbitrary levels.
data LiftedN = LiftedN Nat deriving (Show)

-- | 'Lifted' defines an FKL dialect in which primitives and operators
-- occur either unlifted or lifted once.
data Lifted = Lifted | NotLifted deriving (Show)

-- | 'FExpr' is the target language of the flattening transformation.
data ExprTempl l e = Table Type String L.BaseTableSchema
                   | PApp1 Type Prim1 l (ExprTempl l e)
                   | PApp2 Type Prim2 l (ExprTempl l e) (ExprTempl l e)
                   | PApp3 Type Prim3 l (ExprTempl l e) (ExprTempl l e) (ExprTempl l e)
                   | If Type (ExprTempl l e) (ExprTempl l e) (ExprTempl l e)
                   | BinOp Type L.ScalarBinOp l (ExprTempl l e) (ExprTempl l e)
                   | UnOp Type L.ScalarUnOp l (ExprTempl l e)
                   | Const Type L.Val
                   | Ext e
                   | Let Type L.Ident (ExprTempl l e) (ExprTempl l e)
                   | Var Type L.Ident
                   | MkTuple Type l [(ExprTempl l e)]

data BroadcastExt = Broadcast Nat Type LExpr LExpr

data ShapeExt = Forget Nat Type FExpr
              | Imprint Nat Type FExpr FExpr

type FExpr = ExprTempl Lifted ShapeExt
type LExpr = ExprTempl LiftedN BroadcastExt

-- | Forget does not unsegment the vector. That is: the descriptor
-- might not be normalized and segment descriptors other than 1 might
-- occur. This is propably ok when we know that a concated vector will
-- be unconcated again. We know this statically when introducing
-- concat/unconcat for higher-lifted primitives.

data Prim1 = Length
           | Concat
           | TupElem TupleIndex
           | Sum
           | Avg
           | Minimum
           | Maximum
           | Reverse
           | And
           | Or
           | Nub
           | Number
           | Sort
           | Restrict
           | Group
           | Singleton
           | Only
    deriving (Show, Eq)

data Prim2 = Append
           | Zip
           | CartProduct
           | NestProduct
           | ThetaJoin (L.JoinPredicate L.JoinExpr)
           | NestJoin (L.JoinPredicate L.JoinExpr)
           | SemiJoin (L.JoinPredicate L.JoinExpr)
           | AntiJoin (L.JoinPredicate L.JoinExpr)
           | Dist
           deriving (Show, Eq)

isJoinOp :: Prim2 -> Bool
isJoinOp op =
    case op of
        CartProduct -> True
        NestProduct -> True
        ThetaJoin{} -> True
        NestJoin{}  -> True
        SemiJoin{}  -> True
        AntiJoin{}  -> True
        Append      -> False
        Zip         -> False
        Dist        -> False

data Prim3 = Combine
    deriving (Show, Eq)

instance Typed e => Typed (ExprTempl l e) where
    typeOf (Var t _)           = t
    typeOf (Let t _ _ _)       = t
    typeOf (Table t _ _)       = t
    typeOf (PApp1 t _ _ _)     = t
    typeOf (PApp2 t _ _ _ _)   = t
    typeOf (PApp3 t _ _ _ _ _) = t
    typeOf (If t _ _ _)        = t
    typeOf (BinOp t _ _ _ _)   = t
    typeOf (UnOp t _ _ _)      = t
    typeOf (Const t _)         = t
    typeOf (MkTuple t _ _)     = t
    typeOf (Ext o)             = typeOf o

instance Typed BroadcastExt where
    typeOf (Broadcast _ t _ _) = t

instance Typed ShapeExt where
    typeOf (Forget _ t _)    = t
    typeOf (Imprint _ t _ _) = t

--------------------------------------------------------------------------------
-- Pretty-printing of FKL dialects

superscript :: Int -> Doc
superscript 1 = char '¹'
superscript 2 = char '²'
superscript 3 = char '³'
superscript 4 = char '⁴'
superscript 5 = char '⁵'
superscript 6 = char '⁶'
superscript n = char '^' <> int n

subscript :: Int -> Doc
subscript 1 = char '₁'
subscript 2 = char '₂'
subscript 3 = char '₃'
subscript 4 = char '₄'
subscript 5 = char '₅'
subscript 6 = char '₆'
subscript n = char '^' <> int n

instance Pretty Lifted where
    pretty Lifted    = super $ text "ᴸ"
    pretty NotLifted = empty

instance Pretty LiftedN where
    pretty (LiftedN Zero) = empty
    pretty (LiftedN n)    = super $ superscript (intFromNat n)

instance Pretty Prim1 where
    pretty Length       = combinator $ text "length"
    pretty Concat       = combinator $ text "concat"
    pretty Sum          = combinator $ text "sum"
    pretty Avg          = combinator $ text "avg"<commit_msg>Uniform display of lifting indicators.<commit_after>{-# LANGUAGE FlexibleInstances  #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell    #-}

module Database.DSH.FKL.Lang where

import           Prelude hiding ((<$>))
import           Text.PrettyPrint.ANSI.Leijen
import           Text.Printf

import           Database.DSH.Common.Impossible
import           Database.DSH.Common.Pretty
import           Database.DSH.Common.Nat
import qualified Database.DSH.Common.Lang   as L
import           Database.DSH.Common.Type   (Type, Typed, typeOf)


-- | 'LiftedN' defines an FKL dialect in which primitives and
-- operators might be lifted to arbitrary levels.
data LiftedN = LiftedN Nat deriving (Show)

-- | 'Lifted' defines an FKL dialect in which primitives and operators
-- occur either unlifted or lifted once.
data Lifted = Lifted | NotLifted deriving (Show)

-- | 'FExpr' is the target language of the flattening transformation.
data ExprTempl l e = Table Type String L.BaseTableSchema
                   | PApp1 Type Prim1 l (ExprTempl l e)
                   | PApp2 Type Prim2 l (ExprTempl l e) (ExprTempl l e)
                   | PApp3 Type Prim3 l (ExprTempl l e) (ExprTempl l e) (ExprTempl l e)
                   | If Type (ExprTempl l e) (ExprTempl l e) (ExprTempl l e)
                   | BinOp Type L.ScalarBinOp l (ExprTempl l e) (ExprTempl l e)
                   | UnOp Type L.ScalarUnOp l (ExprTempl l e)
                   | Const Type L.Val
                   | Ext e
                   | Let Type L.Ident (ExprTempl l e) (ExprTempl l e)
                   | Var Type L.Ident
                   | MkTuple Type l [(ExprTempl l e)]

data BroadcastExt = Broadcast Nat Type LExpr LExpr

data ShapeExt = Forget Nat Type FExpr
              | Imprint Nat Type FExpr FExpr

type FExpr = ExprTempl Lifted ShapeExt
type LExpr = ExprTempl LiftedN BroadcastExt

-- | Forget does not unsegment the vector. That is: the descriptor
-- might not be normalized and segment descriptors other than 1 might
-- occur. This is propably ok when we know that a concated vector will
-- be unconcated again. We know this statically when introducing
-- concat/unconcat for higher-lifted primitives.

data Prim1 = Length
           | Concat
           | TupElem TupleIndex
           | Sum
           | Avg
           | Minimum
           | Maximum
           | Reverse
           | And
           | Or
           | Nub
           | Number
           | Sort
           | Restrict
           | Group
           | Singleton
           | Only
    deriving (Show, Eq)

data Prim2 = Append
           | Zip
           | CartProduct
           | NestProduct
           | ThetaJoin (L.JoinPredicate L.JoinExpr)
           | NestJoin (L.JoinPredicate L.JoinExpr)
           | SemiJoin (L.JoinPredicate L.JoinExpr)
           | AntiJoin (L.JoinPredicate L.JoinExpr)
           | Dist
           deriving (Show, Eq)

isJoinOp :: Prim2 -> Bool
isJoinOp op =
    case op of
        CartProduct -> True
        NestProduct -> True
        ThetaJoin{} -> True
        NestJoin{}  -> True
        SemiJoin{}  -> True
        AntiJoin{}  -> True
        Append      -> False
        Zip         -> False
        Dist        -> False

data Prim3 = Combine
    deriving (Show, Eq)

instance Typed e => Typed (ExprTempl l e) where
    typeOf (Var t _)           = t
    typeOf (Let t _ _ _)       = t
    typeOf (Table t _ _)       = t
    typeOf (PApp1 t _ _ _)     = t
    typeOf (PApp2 t _ _ _ _)   = t
    typeOf (PApp3 t _ _ _ _ _) = t
    typeOf (If t _ _ _)        = t
    typeOf (BinOp t _ _ _ _)   = t
    typeOf (UnOp t _ _ _)      = t
    typeOf (Const t _)         = t
    typeOf (MkTuple t _ _)     = t
    typeOf (Ext o)             = typeOf o

instance Typed BroadcastExt where
    typeOf (Broadcast _ t _ _) = t

instance Typed ShapeExt where
    typeOf (Forget _ t _)    = t
    typeOf (Imprint _ t _ _) = t

--------------------------------------------------------------------------------
-- Pretty-printing of FKL dialects

superscript :: Int -> Doc
superscript 0 = char '⁰'
superscript 1 = char '¹'
superscript 2 = char '²'
superscript 3 = char '³'
superscript 4 = char '⁴'
superscript 5 = char '⁵'
superscript 6 = char '⁶'
superscript 7 = char '⁷'
superscript 8 = char '⁸'
superscript n = char '^' <> int n

subscript :: Int -> Doc
subscript 1 = char '₁'
subscript 2 = char '₂'
subscript 3 = char '₃'
subscript 4 = char '₄'
subscript 5 = char '₅'
subscript 6 = char '₆'
subscript 7 = char '₇'
subscript 8 = char '₈'
subscript n = char '_' <> int n

instance Pretty Lifted where
    pretty Lifted    = super $ text "¹"
    pretty NotLifted = super $ text "⁰"

instance Pretty LiftedN where
    pretty (LiftedN n)    = super $ superscript (intFromNat n)

instance Pretty Prim1 where
    pretty Length       = combinator $ text "length"
    pretty Concat       = combinator $ text "concat"
    pretty Sum          = combinator $ text "sum"
    pretty Avg          = combinator $ text "avg"
{'fraction_non_alphanumeric': 0.0940378580828019, 'fraction_numerical': 0.005972264399230692, 'mean_word_length': 1.7807486631016043, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 35, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1049916', 'n_tokens_mistral': 3187, 'n_tokens_neox': 3044, 'n_words': 1067}