{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE NoFieldSelectors #-}

module DataFrame.Internal.Expression where

import qualified Data.Map.Strict as M
import Data.Maybe (fromMaybe)
import Data.String
import qualified Data.Text as T
import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))
import qualified Data.Vector.Generic as VG
import DataFrame.Internal.Column
import qualified DataFrame.Internal.Pretty as P
import Type.Reflection (Typeable, typeOf, typeRep)

{- | Operators are an open typeclass: built-ins get their own 'Typeable' type so the
simplifier can match them by 'cast', and users can add instances. The generic
'UnUDF'/'BinUDF' carriers cover UDFs, dynamic-named, and arithmetic ops.
-}
class (Typeable op) => UnaryOp op where
    unaryFn :: op a b -> a -> b
    unaryName :: op a b -> T.Text
    unarySymbol :: op a b -> Maybe T.Text
    unarySymbol op a b
_ = Maybe Text
forall a. Maybe a
Nothing

class (Typeable op) => BinaryOp op where
    binaryFn :: op a b c -> a -> b -> c
    binaryName :: op a b c -> T.Text
    binarySymbol :: op a b c -> Maybe T.Text
    binarySymbol op a b c
_ = Maybe Text
forall a. Maybe a
Nothing
    binaryCommutative :: op a b c -> Bool
    binaryCommutative op a b c
_ = Bool
False
    binaryPrecedence :: op a b c -> Int
    binaryPrecedence op a b c
_ = Int
9

data UnUDF a b = MkUnaryOp
    { forall a b. UnUDF a b -> a -> b
unaryFn :: a -> b
    , forall a b. UnUDF a b -> Text
unaryName :: T.Text
    , forall a b. UnUDF a b -> Maybe Text
unarySymbol :: Maybe T.Text
    }

data BinUDF a b c = MkBinaryOp
    { forall a b c. BinUDF a b c -> a -> b -> c
binaryFn :: a -> b -> c
    , forall a b c. BinUDF a b c -> Text
binaryName :: T.Text
    , forall a b c. BinUDF a b c -> Maybe Text
binarySymbol :: Maybe T.Text
    , forall a b c. BinUDF a b c -> Bool
binaryCommutative :: Bool
    , forall a b c. BinUDF a b c -> Int
binaryPrecedence :: Int
    }

instance UnaryOp UnUDF where
    unaryFn :: forall a b. UnUDF a b -> a -> b
unaryFn (MkUnaryOp{unaryFn :: forall a b. UnUDF a b -> a -> b
unaryFn = a -> b
f}) = a -> b
f
    unaryName :: forall a b. UnUDF a b -> Text
unaryName (MkUnaryOp{unaryName :: forall a b. UnUDF a b -> Text
unaryName = Text
n}) = Text
n
    unarySymbol :: forall a b. UnUDF a b -> Maybe Text
unarySymbol (MkUnaryOp{unarySymbol :: forall a b. UnUDF a b -> Maybe Text
unarySymbol = Maybe Text
s}) = Maybe Text
s

instance BinaryOp BinUDF where
    binaryFn :: forall a b c. BinUDF a b c -> a -> b -> c
binaryFn (MkBinaryOp{binaryFn :: forall a b c. BinUDF a b c -> a -> b -> c
binaryFn = a -> b -> c
f}) = a -> b -> c
f
    binaryName :: forall a b c. BinUDF a b c -> Text
binaryName (MkBinaryOp{binaryName :: forall a b c. BinUDF a b c -> Text
binaryName = Text
n}) = Text
n
    binarySymbol :: forall a b c. BinUDF a b c -> Maybe Text
binarySymbol (MkBinaryOp{binarySymbol :: forall a b c. BinUDF a b c -> Maybe Text
binarySymbol = Maybe Text
s}) = Maybe Text
s
    binaryCommutative :: forall a b c. BinUDF a b c -> Bool
binaryCommutative (MkBinaryOp{binaryCommutative :: forall a b c. BinUDF a b c -> Bool
binaryCommutative = Bool
c}) = Bool
c
    binaryPrecedence :: forall a b c. BinUDF a b c -> Int
binaryPrecedence (MkBinaryOp{binaryPrecedence :: forall a b c. BinUDF a b c -> Int
binaryPrecedence = Int
p}) = Int
p

data MeanAcc = MeanAcc {-# UNPACK #-} !Double {-# UNPACK #-} !Int
    deriving (Int -> MeanAcc -> ShowS
[MeanAcc] -> ShowS
MeanAcc -> String
(Int -> MeanAcc -> ShowS)
-> (MeanAcc -> String) -> ([MeanAcc] -> ShowS) -> Show MeanAcc
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> MeanAcc -> ShowS
showsPrec :: Int -> MeanAcc -> ShowS
$cshow :: MeanAcc -> String
show :: MeanAcc -> String
$cshowList :: [MeanAcc] -> ShowS
showList :: [MeanAcc] -> ShowS
Show, MeanAcc -> MeanAcc -> Bool
(MeanAcc -> MeanAcc -> Bool)
-> (MeanAcc -> MeanAcc -> Bool) -> Eq MeanAcc
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: MeanAcc -> MeanAcc -> Bool
== :: MeanAcc -> MeanAcc -> Bool
$c/= :: MeanAcc -> MeanAcc -> Bool
/= :: MeanAcc -> MeanAcc -> Bool
Eq, Eq MeanAcc
Eq MeanAcc =>
(MeanAcc -> MeanAcc -> Ordering)
-> (MeanAcc -> MeanAcc -> Bool)
-> (MeanAcc -> MeanAcc -> Bool)
-> (MeanAcc -> MeanAcc -> Bool)
-> (MeanAcc -> MeanAcc -> Bool)
-> (MeanAcc -> MeanAcc -> MeanAcc)
-> (MeanAcc -> MeanAcc -> MeanAcc)
-> Ord MeanAcc
MeanAcc -> MeanAcc -> Bool
MeanAcc -> MeanAcc -> Ordering
MeanAcc -> MeanAcc -> MeanAcc
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 :: MeanAcc -> MeanAcc -> Ordering
compare :: MeanAcc -> MeanAcc -> Ordering
$c< :: MeanAcc -> MeanAcc -> Bool
< :: MeanAcc -> MeanAcc -> Bool
$c<= :: MeanAcc -> MeanAcc -> Bool
<= :: MeanAcc -> MeanAcc -> Bool
$c> :: MeanAcc -> MeanAcc -> Bool
> :: MeanAcc -> MeanAcc -> Bool
$c>= :: MeanAcc -> MeanAcc -> Bool
>= :: MeanAcc -> MeanAcc -> Bool
$cmax :: MeanAcc -> MeanAcc -> MeanAcc
max :: MeanAcc -> MeanAcc -> MeanAcc
$cmin :: MeanAcc -> MeanAcc -> MeanAcc
min :: MeanAcc -> MeanAcc -> MeanAcc
Ord, ReadPrec [MeanAcc]
ReadPrec MeanAcc
Int -> ReadS MeanAcc
ReadS [MeanAcc]
(Int -> ReadS MeanAcc)
-> ReadS [MeanAcc]
-> ReadPrec MeanAcc
-> ReadPrec [MeanAcc]
-> Read MeanAcc
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS MeanAcc
readsPrec :: Int -> ReadS MeanAcc
$creadList :: ReadS [MeanAcc]
readList :: ReadS [MeanAcc]
$creadPrec :: ReadPrec MeanAcc
readPrec :: ReadPrec MeanAcc
$creadListPrec :: ReadPrec [MeanAcc]
readListPrec :: ReadPrec [MeanAcc]
Read)

data AggStrategy a b where
    CollectAgg ::
        (VG.Vector v b, Typeable v) => T.Text -> (v b -> a) -> AggStrategy a b
    FoldAgg :: T.Text -> Maybe a -> (a -> b -> a) -> AggStrategy a b
    MergeAgg ::
        (Columnable acc) =>
        T.Text ->
        acc ->
        (acc -> b -> acc) ->
        (acc -> acc -> acc) ->
        (acc -> a) ->
        AggStrategy a b

data Expr a where
    Col :: (Columnable a) => T.Text -> Expr a
    CastWith ::
        (Columnable a, Columnable b, Read a) =>
        T.Text ->
        T.Text ->
        (Either String a -> b) ->
        Expr b
    CastExprWith ::
        (Columnable a, Columnable b, Columnable src, Read a) =>
        T.Text ->
        (Either String a -> b) ->
        Expr src ->
        Expr b
    Lit :: (Columnable a) => a -> Expr a
    Unary ::
        (UnaryOp op, Columnable a, Columnable b) => op b a -> Expr b -> Expr a
    Binary ::
        (BinaryOp op, Columnable c, Columnable b, Columnable a) =>
        op c b a -> Expr c -> Expr b -> Expr a
    If :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a
    Agg :: (Columnable a, Columnable b) => AggStrategy a b -> Expr b -> Expr a
    Over :: (Columnable a) => [T.Text] -> Expr a -> Expr a

data UExpr where
    UExpr :: (Columnable a) => Expr a -> UExpr

instance Show UExpr where
    show :: UExpr -> String
    show :: UExpr -> String
show (UExpr Expr a
expr) = Expr a -> String
forall a. Show a => a -> String
show Expr a
expr

toUExpr :: (Columnable a) => Expr a -> UExpr
toUExpr :: forall a. Columnable a => Expr a -> UExpr
toUExpr = Expr a -> UExpr
forall a. Columnable a => Expr a -> UExpr
UExpr

fromUExpr :: forall a. (Columnable a) => UExpr -> Maybe (Expr a)
fromUExpr :: forall a. Columnable a => UExpr -> Maybe (Expr a)
fromUExpr (UExpr (Expr a
expr :: Expr b)) = do
    a :~: a
Refl <- TypeRep a -> TypeRep a -> Maybe (a :~: a)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b)
    Expr a -> Maybe (Expr a)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Expr a
Expr a
expr

type NamedExpr = (T.Text, UExpr)

toNamedExpr :: (Columnable a) => T.Text -> Expr a -> NamedExpr
toNamedExpr :: forall a. Columnable a => Text -> Expr a -> NamedExpr
toNamedExpr Text
exprName Expr a
expr = (Text
exprName, Expr a -> UExpr
forall a. Columnable a => Expr a -> UExpr
UExpr Expr a
expr)

instance (Num a, Columnable a) => Num (Expr a) where
    (+) :: Expr a -> Expr a -> Expr a
    + :: Expr a -> Expr a -> Expr a
(+) =
        BinUDF a a a -> Expr a -> Expr a -> Expr a
forall (b :: * -> * -> * -> *) src b a.
(BinaryOp b, Columnable src, Columnable b, Columnable a) =>
b src b a -> Expr src -> Expr b -> Expr a
Binary
            ( MkBinaryOp
                { binaryFn :: a -> a -> a
binaryFn = a -> a -> a
forall a. Num a => a -> a -> a
(+)
                , binaryName :: Text
binaryName = Text
"add"
                , binarySymbol :: Maybe Text
binarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"+"
                , binaryCommutative :: Bool
binaryCommutative = Bool
True
                , binaryPrecedence :: Int
binaryPrecedence = Int
6
                }
            )

    (-) :: Expr a -> Expr a -> Expr a
    (-) =
        BinUDF a a a -> Expr a -> Expr a -> Expr a
forall (b :: * -> * -> * -> *) src b a.
(BinaryOp b, Columnable src, Columnable b, Columnable a) =>
b src b a -> Expr src -> Expr b -> Expr a
Binary
            ( MkBinaryOp
                { binaryFn :: a -> a -> a
binaryFn = (-)
                , binaryName :: Text
binaryName = Text
"sub"
                , binarySymbol :: Maybe Text
binarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"-"
                , binaryCommutative :: Bool
binaryCommutative = Bool
False
                , binaryPrecedence :: Int
binaryPrecedence = Int
6
                }
            )

    (*) :: Expr a -> Expr a -> Expr a
    * :: Expr a -> Expr a -> Expr a
(*) =
        BinUDF a a a -> Expr a -> Expr a -> Expr a
forall (b :: * -> * -> * -> *) src b a.
(BinaryOp b, Columnable src, Columnable b, Columnable a) =>
b src b a -> Expr src -> Expr b -> Expr a
Binary
            ( MkBinaryOp
                { binaryFn :: a -> a -> a
binaryFn = a -> a -> a
forall a. Num a => a -> a -> a
(*)
                , binaryName :: Text
binaryName = Text
"mult"
                , binarySymbol :: Maybe Text
binarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"*"
                , binaryCommutative :: Bool
binaryCommutative = Bool
True
                , binaryPrecedence :: Int
binaryPrecedence = Int
7
                }
            )

    fromInteger :: Integer -> Expr a
    fromInteger :: Integer -> Expr a
fromInteger = a -> Expr a
forall a. Columnable a => a -> Expr a
Lit (a -> Expr a) -> (Integer -> a) -> Integer -> Expr a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> a
forall a. Num a => Integer -> a
fromInteger

    negate :: Expr a -> Expr a
    negate :: Expr a -> Expr a
negate =
        UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary
            (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Num a => a -> a
negate, unaryName :: Text
unaryName = Text
"negate", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})

    abs :: (Num a) => Expr a -> Expr a
    abs :: Num a => Expr a -> Expr a
abs = UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Num a => a -> a
abs, unaryName :: Text
unaryName = Text
"abs", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})

    signum :: (Num a) => Expr a -> Expr a
    signum :: Num a => Expr a -> Expr a
signum =
        UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary
            (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Num a => a -> a
signum, unaryName :: Text
unaryName = Text
"signum", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})

add :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a
add :: forall a. (Num a, Columnable a) => Expr a -> Expr a -> Expr a
add = Expr a -> Expr a -> Expr a
forall a. Num a => a -> a -> a
(+)

sub :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a
sub :: forall a. (Num a, Columnable a) => Expr a -> Expr a -> Expr a
sub = (-)

mult :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a
mult :: forall a. (Num a, Columnable a) => Expr a -> Expr a -> Expr a
mult = Expr a -> Expr a -> Expr a
forall a. Num a => a -> a -> a
(*)

instance (Fractional a, Columnable a) => Fractional (Expr a) where
    fromRational :: (Fractional a, Columnable a) => Rational -> Expr a
    fromRational :: (Fractional a, Columnable a) => Rational -> Expr a
fromRational = a -> Expr a
forall a. Columnable a => a -> Expr a
Lit (a -> Expr a) -> (Rational -> a) -> Rational -> Expr a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rational -> a
forall a. Fractional a => Rational -> a
fromRational

    (/) :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a
    / :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a
(/) =
        BinUDF a a a -> Expr a -> Expr a -> Expr a
forall (b :: * -> * -> * -> *) src b a.
(BinaryOp b, Columnable src, Columnable b, Columnable a) =>
b src b a -> Expr src -> Expr b -> Expr a
Binary
            ( MkBinaryOp
                { binaryFn :: a -> a -> a
binaryFn = a -> a -> a
forall a. Fractional a => a -> a -> a
(/)
                , binaryName :: Text
binaryName = Text
"divide"
                , binarySymbol :: Maybe Text
binarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"/"
                , binaryCommutative :: Bool
binaryCommutative = Bool
False
                , binaryPrecedence :: Int
binaryPrecedence = Int
7
                }
            )

divide :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a
divide :: forall a.
(Fractional a, Columnable a) =>
Expr a -> Expr a -> Expr a
divide = Expr a -> Expr a -> Expr a
forall a. Fractional a => a -> a -> a
(/)

instance (IsString a, Columnable a) => IsString (Expr a) where
    fromString :: String -> Expr a
    fromString :: String -> Expr a
fromString String
s = a -> Expr a
forall a. Columnable a => a -> Expr a
Lit (String -> a
forall a. IsString a => String -> a
fromString String
s)

instance (Floating a, Columnable a) => Floating (Expr a) where
    pi :: (Floating a, Columnable a) => Expr a
    pi :: (Floating a, Columnable a) => Expr a
pi = a -> Expr a
forall a. Columnable a => a -> Expr a
Lit a
forall a. Floating a => a
pi
    exp :: (Floating a, Columnable a) => Expr a -> Expr a
    exp :: (Floating a, Columnable a) => Expr a -> Expr a
exp = UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Floating a => a -> a
exp, unaryName :: Text
unaryName = Text
"exp", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})
    sqrt :: (Floating a, Columnable a) => Expr a -> Expr a
    sqrt :: (Floating a, Columnable a) => Expr a -> Expr a
sqrt =
        UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Floating a => a -> a
sqrt, unaryName :: Text
unaryName = Text
"sqrt", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})
    (**) :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a
    ** :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a
(**) =
        BinUDF a a a -> Expr a -> Expr a -> Expr a
forall (b :: * -> * -> * -> *) src b a.
(BinaryOp b, Columnable src, Columnable b, Columnable a) =>
b src b a -> Expr src -> Expr b -> Expr a
Binary
            ( MkBinaryOp
                { binaryFn :: a -> a -> a
binaryFn = a -> a -> a
forall a. Floating a => a -> a -> a
(**)
                , binaryName :: Text
binaryName = Text
"exponentiate"
                , binarySymbol :: Maybe Text
binarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"**"
                , binaryCommutative :: Bool
binaryCommutative = Bool
False
                , binaryPrecedence :: Int
binaryPrecedence = Int
8
                }
            )
    log :: (Floating a, Columnable a) => Expr a -> Expr a
    log :: (Floating a, Columnable a) => Expr a -> Expr a
log = UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Floating a => a -> a
log, unaryName :: Text
unaryName = Text
"log", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})
    logBase :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a
    logBase :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a
logBase =
        BinUDF a a a -> Expr a -> Expr a -> Expr a
forall (b :: * -> * -> * -> *) src b a.
(BinaryOp b, Columnable src, Columnable b, Columnable a) =>
b src b a -> Expr src -> Expr b -> Expr a
Binary
            ( MkBinaryOp
                { binaryFn :: a -> a -> a
binaryFn = a -> a -> a
forall a. Floating a => a -> a -> a
logBase
                , binaryName :: Text
binaryName = Text
"logBase"
                , binarySymbol :: Maybe Text
binarySymbol = Maybe Text
forall a. Maybe a
Nothing
                , binaryCommutative :: Bool
binaryCommutative = Bool
False
                , binaryPrecedence :: Int
binaryPrecedence = Int
1
                }
            )
    sin :: (Floating a, Columnable a) => Expr a -> Expr a
    sin :: (Floating a, Columnable a) => Expr a -> Expr a
sin = UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Floating a => a -> a
sin, unaryName :: Text
unaryName = Text
"sin", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})
    cos :: (Floating a, Columnable a) => Expr a -> Expr a
    cos :: (Floating a, Columnable a) => Expr a -> Expr a
cos = UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Floating a => a -> a
cos, unaryName :: Text
unaryName = Text
"cos", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})
    tan :: (Floating a, Columnable a) => Expr a -> Expr a
    tan :: (Floating a, Columnable a) => Expr a -> Expr a
tan = UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Floating a => a -> a
tan, unaryName :: Text
unaryName = Text
"tan", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})
    asin :: (Floating a, Columnable a) => Expr a -> Expr a
    asin :: (Floating a, Columnable a) => Expr a -> Expr a
asin =
        UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Floating a => a -> a
asin, unaryName :: Text
unaryName = Text
"asin", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})
    acos :: (Floating a, Columnable a) => Expr a -> Expr a
    acos :: (Floating a, Columnable a) => Expr a -> Expr a
acos =
        UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Floating a => a -> a
acos, unaryName :: Text
unaryName = Text
"acos", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})
    atan :: (Floating a, Columnable a) => Expr a -> Expr a
    atan :: (Floating a, Columnable a) => Expr a -> Expr a
atan =
        UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Floating a => a -> a
atan, unaryName :: Text
unaryName = Text
"atan", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})
    sinh :: (Floating a, Columnable a) => Expr a -> Expr a
    sinh :: (Floating a, Columnable a) => Expr a -> Expr a
sinh =
        UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Floating a => a -> a
sinh, unaryName :: Text
unaryName = Text
"sinh", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})
    cosh :: (Floating a, Columnable a) => Expr a -> Expr a
    cosh :: (Floating a, Columnable a) => Expr a -> Expr a
cosh =
        UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Floating a => a -> a
cosh, unaryName :: Text
unaryName = Text
"cosh", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})
    asinh :: (Floating a, Columnable a) => Expr a -> Expr a
    asinh :: (Floating a, Columnable a) => Expr a -> Expr a
asinh =
        UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary
            (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Floating a => a -> a
asinh, unaryName :: Text
unaryName = Text
"asinh", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})
    acosh :: (Floating a, Columnable a) => Expr a -> Expr a
    acosh :: (Floating a, Columnable a) => Expr a -> Expr a
acosh =
        UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary
            (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Floating a => a -> a
acosh, unaryName :: Text
unaryName = Text
"acosh", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})
    atanh :: (Floating a, Columnable a) => Expr a -> Expr a
    atanh :: (Floating a, Columnable a) => Expr a -> Expr a
atanh =
        UnUDF a a -> Expr a -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary
            (MkUnaryOp{unaryFn :: a -> a
unaryFn = a -> a
forall a. Floating a => a -> a
atanh, unaryName :: Text
unaryName = Text
"atanh", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})

instance (Show a) => Show (Expr a) where
    show :: Expr a -> String
    show :: Expr a -> String
show (Col Text
name) = String
"(col @" String -> ShowS
forall a. [a] -> [a] -> [a]
++ TypeRep a -> String
forall a. Show a => a -> String
show (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
forall a. Show a => a -> String
show Text
name String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    show (CastWith Text
name Text
tag Either String a -> a
_) = String
"(castWith " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
forall a. Show a => a -> String
show Text
tag String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
forall a. Show a => a -> String
show Text
name String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    show (CastExprWith Text
tag Either String a -> a
_ Expr src
inner) = String
"(castExprWith " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
forall a. Show a => a -> String
show Text
tag String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr src -> String
forall a. Show a => a -> String
show Expr src
inner String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    show (Lit a
value) = String
"(lit (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
value String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"))"
    show (If Expr Bool
cond Expr a
l Expr a
r) = String
"(ifThenElse " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr Bool -> String
forall a. Show a => a -> String
show Expr Bool
cond String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr a -> String
forall a. Show a => a -> String
show Expr a
l String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr a -> String
forall a. Show a => a -> String
show Expr a
r String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    show (Unary op b a
op Expr b
value) = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack (op b a -> Text
forall a b. op a b -> Text
forall (op :: * -> * -> *) a b. UnaryOp op => op a b -> Text
unaryName op b a
op) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr b -> String
forall a. Show a => a -> String
show Expr b
value String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    show (Binary op c b a
op Expr c
a Expr b
b) = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack (op c b a -> Text
forall a b c. op a b c -> Text
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> Text
binaryName op c b a
op) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr c -> String
forall a. Show a => a -> String
show Expr c
a String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr b -> String
forall a. Show a => a -> String
show Expr b
b String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    show (Agg (CollectAgg Text
op v b -> a
_) Expr b
expr) = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack Text
op String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr b -> String
forall a. Show a => a -> String
show Expr b
expr String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    show (Agg (FoldAgg Text
op Maybe a
_ a -> b -> a
_) Expr b
expr) = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack Text
op String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr b -> String
forall a. Show a => a -> String
show Expr b
expr String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    show (Agg (MergeAgg Text
op acc
_ acc -> b -> acc
_ acc -> acc -> acc
_ acc -> a
_) Expr b
expr) = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack Text
op String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr b -> String
forall a. Show a => a -> String
show Expr b
expr String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    show (Over [Text]
keys Expr a
inner) = String
"(over " String -> ShowS
forall a. [a] -> [a] -> [a]
++ [Text] -> String
forall a. Show a => a -> String
show [Text]
keys String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr a -> String
forall a. Show a => a -> String
show Expr a
inner String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

normalize :: (Show a, Typeable a) => Expr a -> Expr a
normalize :: forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr a
expr = case Expr a
expr of
    Col Text
name -> Text -> Expr a
forall a. Columnable a => Text -> Expr a
Col Text
name
    CastWith Text
n Text
t Either String a -> a
f -> Text -> Text -> (Either String a -> a) -> Expr a
forall b b.
(Columnable b, Columnable b, Read b) =>
Text -> Text -> (Either String b -> b) -> Expr b
CastWith Text
n Text
t Either String a -> a
f
    CastExprWith Text
t Either String a -> a
f Expr src
e -> Text -> (Either String a -> a) -> Expr src -> Expr a
forall b b src.
(Columnable b, Columnable b, Columnable src, Read b) =>
Text -> (Either String b -> b) -> Expr src -> Expr b
CastExprWith Text
t Either String a -> a
f (Expr src -> Expr src
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr src
e)
    Lit a
val -> a -> Expr a
forall a. Columnable a => a -> Expr a
Lit a
val
    If Expr Bool
cond Expr a
th Expr a
el -> Expr Bool -> Expr a -> Expr a -> Expr a
forall a. Columnable a => Expr Bool -> Expr a -> Expr a -> Expr a
If (Expr Bool -> Expr Bool
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr Bool
cond) (Expr a -> Expr a
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr a
th) (Expr a -> Expr a
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr a
el)
    Unary op b a
op Expr b
e -> op b a -> Expr b -> Expr a
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary op b a
op (Expr b -> Expr b
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr b
e)
    Binary op c b a
op Expr c
e1 Expr b
e2
        | op c b a -> Bool
forall a b c. op a b c -> Bool
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> Bool
binaryCommutative op c b a
op ->
            let n1 :: Expr c
n1 = Expr c -> Expr c
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr c
e1
                n2 :: Expr b
n2 = Expr b -> Expr b
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr b
e2
             in case TypeRep (Expr c) -> TypeRep (Expr b) -> Maybe (Expr c :~: Expr b)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (Expr c -> TypeRep (Expr c)
forall a. Typeable a => a -> TypeRep a
typeOf Expr c
n1) (Expr b -> TypeRep (Expr b)
forall a. Typeable a => a -> TypeRep a
typeOf Expr b
n2) of
                    Maybe (Expr c :~: Expr b)
Nothing -> Expr a
expr
                    Just Expr c :~: Expr b
Refl ->
                        if Expr c -> Expr c -> Ordering
forall a. Expr a -> Expr a -> Ordering
compareExpr Expr c
n1 Expr c
Expr b
n2 Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
GT
                            then op c b a -> Expr c -> Expr b -> Expr a
forall (b :: * -> * -> * -> *) src b a.
(BinaryOp b, Columnable src, Columnable b, Columnable a) =>
b src b a -> Expr src -> Expr b -> Expr a
Binary op c b a
op Expr c
Expr b
n2 Expr c
Expr b
n1
                            else op c b a -> Expr c -> Expr b -> Expr a
forall (b :: * -> * -> * -> *) src b a.
(BinaryOp b, Columnable src, Columnable b, Columnable a) =>
b src b a -> Expr src -> Expr b -> Expr a
Binary op c b a
op Expr c
n1 Expr b
n2
        | Bool
otherwise -> op c b a -> Expr c -> Expr b -> Expr a
forall (b :: * -> * -> * -> *) src b a.
(BinaryOp b, Columnable src, Columnable b, Columnable a) =>
b src b a -> Expr src -> Expr b -> Expr a
Binary op c b a
op (Expr c -> Expr c
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr c
e1) (Expr b -> Expr b
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr b
e2)
    Agg AggStrategy a b
strat Expr b
e -> AggStrategy a b -> Expr b -> Expr a
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg AggStrategy a b
strat (Expr b -> Expr b
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr b
e)
    Over [Text]
keys Expr a
inner -> [Text] -> Expr a -> Expr a
forall a. Columnable a => [Text] -> Expr a -> Expr a
Over [Text]
keys (Expr a -> Expr a
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr a
inner)

-- Compare expressions for ordering (used in normalization)
compareExpr :: Expr a -> Expr a -> Ordering
compareExpr :: forall a. Expr a -> Expr a -> Ordering
compareExpr Expr a
e1 Expr a
e2 = String -> String -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Expr a -> String
forall a. Expr a -> String
exprKey Expr a
e1) (Expr a -> String
forall a. Expr a -> String
exprKey Expr a
e2)
  where
    exprKey :: Expr a -> String
    exprKey :: forall a. Expr a -> String
exprKey (Col Text
name) = String
"0:" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack Text
name
    exprKey (CastWith Text
name Text
tag Either String a -> a
_) = String
"0CW:" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack Text
name String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
":" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack Text
tag
    exprKey (CastExprWith Text
tag Either String a -> a
_ Expr src
_) = String
"0CE:" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack Text
tag
    exprKey (Lit a
val) = String
"1:" String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
val
    exprKey (If Expr Bool
c Expr a
t Expr a
e) = String
"2:" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr Bool -> String
forall a. Expr a -> String
exprKey Expr Bool
c String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr a -> String
forall a. Expr a -> String
exprKey Expr a
t String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr a -> String
forall a. Expr a -> String
exprKey Expr a
e
    exprKey (Unary op b a
op Expr b
e) = String
"3:" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack (op b a -> Text
forall a b. op a b -> Text
forall (op :: * -> * -> *) a b. UnaryOp op => op a b -> Text
unaryName op b a
op) String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr b -> String
forall a. Expr a -> String
exprKey Expr b
e
    exprKey (Binary op c b a
op Expr c
e1' Expr b
e2') = String
"4:" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack (op c b a -> Text
forall a b c. op a b c -> Text
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> Text
binaryName op c b a
op) String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr c -> String
forall a. Expr a -> String
exprKey Expr c
e1' String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr b -> String
forall a. Expr a -> String
exprKey Expr b
e2'
    exprKey (Agg (CollectAgg Text
name v b -> a
_) Expr b
e) = String
"5:" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack Text
name String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr b -> String
forall a. Expr a -> String
exprKey Expr b
e
    exprKey (Agg (FoldAgg Text
name Maybe a
_ a -> b -> a
_) Expr b
e) = String
"5:" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack Text
name String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr b -> String
forall a. Expr a -> String
exprKey Expr b
e
    exprKey (Agg (MergeAgg Text
name acc
_ acc -> b -> acc
_ acc -> acc -> acc
_ acc -> a
_) Expr b
e) = String
"5:" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack Text
name String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr b -> String
forall a. Expr a -> String
exprKey Expr b
e
    exprKey (Over [Text]
keys Expr a
e) = String
"6:over:" String -> ShowS
forall a. [a] -> [a] -> [a]
++ [Text] -> String
forall a. Show a => a -> String
show [Text]
keys String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr a -> String
forall a. Expr a -> String
exprKey Expr a
e

eqExpr :: forall a. (Columnable a) => Expr a -> Expr a -> Bool
eqExpr :: forall a. Columnable a => Expr a -> Expr a -> Bool
eqExpr Expr a
l Expr a
r = Expr a -> Expr a -> Bool
eqNormalized (Expr a -> Expr a
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr a
l) (Expr a -> Expr a
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr a
r)
  where
    exprEq :: (Columnable b, Columnable c) => Expr b -> Expr c -> Bool
    exprEq :: forall b c.
(Columnable b, Columnable c) =>
Expr b -> Expr c -> Bool
exprEq Expr b
e1 Expr c
e2 = case TypeRep (Expr b) -> TypeRep (Expr c) -> Maybe (Expr b :~: Expr c)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (Expr b -> TypeRep (Expr b)
forall a. Typeable a => a -> TypeRep a
typeOf Expr b
e1) (Expr c -> TypeRep (Expr c)
forall a. Typeable a => a -> TypeRep a
typeOf Expr c
e2) of
        Just Expr b :~: Expr c
Refl -> Expr b -> Expr b -> Bool
forall a. Columnable a => Expr a -> Expr a -> Bool
eqExpr Expr b
e1 Expr b
Expr c
e2
        Maybe (Expr b :~: Expr c)
Nothing -> Bool
False
    eqNormalized :: Expr a -> Expr a -> Bool
    eqNormalized :: Expr a -> Expr a -> Bool
eqNormalized (Col Text
n1) (Col Text
n2) = Text
n1 Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
n2
    eqNormalized (CastWith Text
n1 Text
t1 Either String a -> a
_) (CastWith Text
n2 Text
t2 Either String a -> a
_) = Text
n1 Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
n2 Bool -> Bool -> Bool
&& Text
t1 Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
t2
    eqNormalized (CastExprWith Text
t1 Either String a -> a
_ Expr src
e1) (CastExprWith Text
t2 Either String a -> a
_ Expr src
e2) = Text
t1 Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
t2 Bool -> Bool -> Bool
&& Expr src
e1 Expr src -> Expr src -> Bool
forall b c.
(Columnable b, Columnable c) =>
Expr b -> Expr c -> Bool
`exprEq` Expr src
e2
    eqNormalized (Lit a
v1) (Lit a
v2) = a
v1 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
v2
    eqNormalized (If Expr Bool
c1 Expr a
t1 Expr a
e1) (If Expr Bool
c2 Expr a
t2 Expr a
e2) =
        Expr Bool -> Expr Bool -> Bool
forall a. Columnable a => Expr a -> Expr a -> Bool
eqExpr Expr Bool
c1 Expr Bool
c2 Bool -> Bool -> Bool
&& Expr a
t1 Expr a -> Expr a -> Bool
forall b c.
(Columnable b, Columnable c) =>
Expr b -> Expr c -> Bool
`exprEq` Expr a
t2 Bool -> Bool -> Bool
&& Expr a
e1 Expr a -> Expr a -> Bool
forall b c.
(Columnable b, Columnable c) =>
Expr b -> Expr c -> Bool
`exprEq` Expr a
e2
    eqNormalized (Unary op b a
op1 Expr b
e1) (Unary op b a
op2 Expr b
e2) = op b a -> Text
forall a b. op a b -> Text
forall (op :: * -> * -> *) a b. UnaryOp op => op a b -> Text
unaryName op b a
op1 Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== op b a -> Text
forall a b. op a b -> Text
forall (op :: * -> * -> *) a b. UnaryOp op => op a b -> Text
unaryName op b a
op2 Bool -> Bool -> Bool
&& Expr b
e1 Expr b -> Expr b -> Bool
forall b c.
(Columnable b, Columnable c) =>
Expr b -> Expr c -> Bool
`exprEq` Expr b
e2
    eqNormalized (Binary op c b a
op1 Expr c
e1a Expr b
e1b) (Binary op c b a
op2 Expr c
e2a Expr b
e2b) = op c b a -> Text
forall a b c. op a b c -> Text
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> Text
binaryName op c b a
op1 Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== op c b a -> Text
forall a b c. op a b c -> Text
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> Text
binaryName op c b a
op2 Bool -> Bool -> Bool
&& Expr c
e1a Expr c -> Expr c -> Bool
forall b c.
(Columnable b, Columnable c) =>
Expr b -> Expr c -> Bool
`exprEq` Expr c
e2a Bool -> Bool -> Bool
&& Expr b
e1b Expr b -> Expr b -> Bool
forall b c.
(Columnable b, Columnable c) =>
Expr b -> Expr c -> Bool
`exprEq` Expr b
e2b
    eqNormalized (Agg (CollectAgg Text
n1 v b -> a
_) Expr b
e1) (Agg (CollectAgg Text
n2 v b -> a
_) Expr b
e2) =
        Text
n1 Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
n2 Bool -> Bool -> Bool
&& Expr b
e1 Expr b -> Expr b -> Bool
forall b c.
(Columnable b, Columnable c) =>
Expr b -> Expr c -> Bool
`exprEq` Expr b
e2
    eqNormalized (Agg (FoldAgg Text
n1 Maybe a
_ a -> b -> a
_) Expr b
e1) (Agg (FoldAgg Text
n2 Maybe a
_ a -> b -> a
_) Expr b
e2) =
        Text
n1 Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
n2 Bool -> Bool -> Bool
&& Expr b
e1 Expr b -> Expr b -> Bool
forall b c.
(Columnable b, Columnable c) =>
Expr b -> Expr c -> Bool
`exprEq` Expr b
e2
    eqNormalized (Agg (MergeAgg Text
n1 acc
_ acc -> b -> acc
_ acc -> acc -> acc
_ acc -> a
_) Expr b
e1) (Agg (MergeAgg Text
n2 acc
_ acc -> b -> acc
_ acc -> acc -> acc
_ acc -> a
_) Expr b
e2) =
        Text
n1 Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
n2 Bool -> Bool -> Bool
&& Expr b
e1 Expr b -> Expr b -> Bool
forall b c.
(Columnable b, Columnable c) =>
Expr b -> Expr c -> Bool
`exprEq` Expr b
e2
    eqNormalized (Over [Text]
k1 Expr a
e1) (Over [Text]
k2 Expr a
e2) = [Text]
k1 [Text] -> [Text] -> Bool
forall a. Eq a => a -> a -> Bool
== [Text]
k2 Bool -> Bool -> Bool
&& Expr a
e1 Expr a -> Expr a -> Bool
forall b c.
(Columnable b, Columnable c) =>
Expr b -> Expr c -> Bool
`exprEq` Expr a
e2
    eqNormalized Expr a
_ Expr a
_ = Bool
False

replaceExpr ::
    forall a b c.
    (Columnable a, Columnable b, Columnable c) =>
    Expr a -> Expr b -> Expr c -> Expr c
replaceExpr :: forall a b c.
(Columnable a, Columnable b, Columnable c) =>
Expr a -> Expr b -> Expr c -> Expr c
replaceExpr Expr a
new Expr b
old Expr c
expr = case TypeRep b -> TypeRep c -> Maybe (b :~: c)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @c) of
    Just b :~: c
Refl -> case TypeRep a -> TypeRep c -> Maybe (a :~: c)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @c) of
        Just a :~: c
Refl -> if Expr b -> Expr b -> Bool
forall a. Columnable a => Expr a -> Expr a -> Bool
eqExpr Expr b
old Expr b
Expr c
expr then Expr a
Expr c
new else Expr c
replace'
        Maybe (a :~: c)
Nothing -> Expr c
expr
    Maybe (b :~: c)
Nothing -> Expr c
replace'
  where
    replace' :: Expr c
replace' = case Expr c
expr of
        (Col Text
_) -> Expr c
expr
        (CastWith{}) -> Expr c
expr
        (CastExprWith Text
t Either String a -> c
f Expr src
e) -> Text -> (Either String a -> c) -> Expr src -> Expr c
forall b b src.
(Columnable b, Columnable b, Columnable src, Read b) =>
Text -> (Either String b -> b) -> Expr src -> Expr b
CastExprWith Text
t Either String a -> c
f (Expr a -> Expr b -> Expr src -> Expr src
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
Expr a -> Expr b -> Expr c -> Expr c
replaceExpr Expr a
new Expr b
old Expr src
e)
        (Lit c
_) -> Expr c
expr
        (If Expr Bool
cond Expr c
l Expr c
r) ->
            Expr Bool -> Expr c -> Expr c -> Expr c
forall a. Columnable a => Expr Bool -> Expr a -> Expr a -> Expr a
If (Expr a -> Expr b -> Expr Bool -> Expr Bool
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
Expr a -> Expr b -> Expr c -> Expr c
replaceExpr Expr a
new Expr b
old Expr Bool
cond) (Expr a -> Expr b -> Expr c -> Expr c
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
Expr a -> Expr b -> Expr c -> Expr c
replaceExpr Expr a
new Expr b
old Expr c
l) (Expr a -> Expr b -> Expr c -> Expr c
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
Expr a -> Expr b -> Expr c -> Expr c
replaceExpr Expr a
new Expr b
old Expr c
r)
        (Unary op b c
op Expr b
value) -> op b c -> Expr b -> Expr c
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary op b c
op (Expr a -> Expr b -> Expr b -> Expr b
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
Expr a -> Expr b -> Expr c -> Expr c
replaceExpr Expr a
new Expr b
old Expr b
value)
        (Binary op c b c
op Expr c
l Expr b
r) -> op c b c -> Expr c -> Expr b -> Expr c
forall (b :: * -> * -> * -> *) src b a.
(BinaryOp b, Columnable src, Columnable b, Columnable a) =>
b src b a -> Expr src -> Expr b -> Expr a
Binary op c b c
op (Expr a -> Expr b -> Expr c -> Expr c
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
Expr a -> Expr b -> Expr c -> Expr c
replaceExpr Expr a
new Expr b
old Expr c
l) (Expr a -> Expr b -> Expr b -> Expr b
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
Expr a -> Expr b -> Expr c -> Expr c
replaceExpr Expr a
new Expr b
old Expr b
r)
        (Agg AggStrategy c b
op Expr b
inner) -> AggStrategy c b -> Expr b -> Expr c
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg AggStrategy c b
op (Expr a -> Expr b -> Expr b -> Expr b
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
Expr a -> Expr b -> Expr c -> Expr c
replaceExpr Expr a
new Expr b
old Expr b
inner)
        (Over [Text]
keys Expr c
inner) -> [Text] -> Expr c -> Expr c
forall a. Columnable a => [Text] -> Expr a -> Expr a
Over [Text]
keys (Expr a -> Expr b -> Expr c -> Expr c
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
Expr a -> Expr b -> Expr c -> Expr c
replaceExpr Expr a
new Expr b
old Expr c
inner)

{- | Simultaneously substitute 'Col' references from a name→expression map in a
single parallel pass, so a swap like @{a ↦ col b, b ↦ col a}@ works. Raw-text
references (in 'CastWith', 'Over' keys) are left untouched; type mismatch raises.
-}
substituteColumns ::
    forall a. (Columnable a) => M.Map T.Text UExpr -> Expr a -> Expr a
substituteColumns :: forall a. Columnable a => Map Text UExpr -> Expr a -> Expr a
substituteColumns Map Text UExpr
subs = Expr a -> Expr a
forall b. Columnable b => Expr b -> Expr b
go
  where
    go :: forall b. (Columnable b) => Expr b -> Expr b
    go :: forall b. Columnable b => Expr b -> Expr b
go e :: Expr b
e@(Col Text
name) = case Text -> Map Text UExpr -> Maybe UExpr
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Text
name Map Text UExpr
subs of
        Maybe UExpr
Nothing -> Expr b
e
        Just (UExpr (Expr a
repl :: Expr c)) -> case TypeRep b -> TypeRep a -> Maybe (b :~: a)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @c) of
            Just b :~: a
Refl -> Expr b
Expr a
repl
            Maybe (b :~: a)
Nothing ->
                String -> Expr b
forall a. HasCallStack => String -> a
error (String -> Expr b) -> String -> Expr b
forall a b. (a -> b) -> a -> b
$
                    String
"substituteColumns: type mismatch for column "
                        String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
forall a. Show a => a -> String
show Text
name
                        String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"; column has type "
                        String -> ShowS
forall a. [a] -> [a] -> [a]
++ TypeRep b -> String
forall a. Show a => a -> String
show (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b)
                        String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" but replacement has type "
                        String -> ShowS
forall a. [a] -> [a] -> [a]
++ TypeRep a -> String
forall a. Show a => a -> String
show (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @c)
    go e :: Expr b
e@(CastWith{}) = Expr b
e
    go (CastExprWith Text
t Either String a -> b
f Expr src
e) = Text -> (Either String a -> b) -> Expr src -> Expr b
forall b b src.
(Columnable b, Columnable b, Columnable src, Read b) =>
Text -> (Either String b -> b) -> Expr src -> Expr b
CastExprWith Text
t Either String a -> b
f (Expr src -> Expr src
forall b. Columnable b => Expr b -> Expr b
go Expr src
e)
    go e :: Expr b
e@(Lit b
_) = Expr b
e
    go (If Expr Bool
cond Expr b
l Expr b
r) = Expr Bool -> Expr b -> Expr b -> Expr b
forall a. Columnable a => Expr Bool -> Expr a -> Expr a -> Expr a
If (Expr Bool -> Expr Bool
forall b. Columnable b => Expr b -> Expr b
go Expr Bool
cond) (Expr b -> Expr b
forall b. Columnable b => Expr b -> Expr b
go Expr b
l) (Expr b -> Expr b
forall b. Columnable b => Expr b -> Expr b
go Expr b
r)
    go (Unary op b b
op Expr b
value) = op b b -> Expr b -> Expr b
forall (b :: * -> * -> *) a src.
(UnaryOp b, Columnable a, Columnable src) =>
b src a -> Expr src -> Expr a
Unary op b b
op (Expr b -> Expr b
forall b. Columnable b => Expr b -> Expr b
go Expr b
value)
    go (Binary op c b b
op Expr c
l Expr b
r) = op c b b -> Expr c -> Expr b -> Expr b
forall (b :: * -> * -> * -> *) src b a.
(BinaryOp b, Columnable src, Columnable b, Columnable a) =>
b src b a -> Expr src -> Expr b -> Expr a
Binary op c b b
op (Expr c -> Expr c
forall b. Columnable b => Expr b -> Expr b
go Expr c
l) (Expr b -> Expr b
forall b. Columnable b => Expr b -> Expr b
go Expr b
r)
    go (Agg AggStrategy b b
op Expr b
inner) = AggStrategy b b -> Expr b -> Expr b
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg AggStrategy b b
op (Expr b -> Expr b
forall b. Columnable b => Expr b -> Expr b
go Expr b
inner)
    go (Over [Text]
keys Expr b
inner) = [Text] -> Expr b -> Expr b
forall a. Columnable a => [Text] -> Expr a -> Expr a
Over [Text]
keys (Expr b -> Expr b
forall b. Columnable b => Expr b -> Expr b
go Expr b
inner)

eSize :: Expr a -> Int
eSize :: forall a. Expr a -> Int
eSize (Col Text
_) = Int
1
eSize (CastWith{}) = Int
1
eSize (CastExprWith Text
_ Either String a -> a
_ Expr src
e) = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr src -> Int
forall a. Expr a -> Int
eSize Expr src
e
eSize (Lit a
_) = Int
1
eSize (If Expr Bool
c Expr a
l Expr a
r) = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr Bool -> Int
forall a. Expr a -> Int
eSize Expr Bool
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr a -> Int
forall a. Expr a -> Int
eSize Expr a
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr a -> Int
forall a. Expr a -> Int
eSize Expr a
r
eSize (Unary op b a
_ Expr b
e) = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr b -> Int
forall a. Expr a -> Int
eSize Expr b
e
eSize (Binary op c b a
_ Expr c
l Expr b
r) = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr c -> Int
forall a. Expr a -> Int
eSize Expr c
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr b -> Int
forall a. Expr a -> Int
eSize Expr b
r
eSize (Agg AggStrategy a b
_strategy Expr b
expr) = Expr b -> Int
forall a. Expr a -> Int
eSize Expr b
expr Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
eSize (Over [Text]
_ Expr a
inner) = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr a -> Int
forall a. Expr a -> Int
eSize Expr a
inner

getColumns :: Expr a -> [T.Text]
getColumns :: forall a. Expr a -> [Text]
getColumns (Col Text
cName) = [Text
cName]
getColumns (CastWith Text
name Text
_ Either String a -> a
_) = [Text
name]
getColumns (CastExprWith Text
_ Either String a -> a
_ Expr src
e) = Expr src -> [Text]
forall a. Expr a -> [Text]
getColumns Expr src
e
getColumns _expr :: Expr a
_expr@(Lit a
_) = []
getColumns (If Expr Bool
cond Expr a
l Expr a
r) = Expr Bool -> [Text]
forall a. Expr a -> [Text]
getColumns Expr Bool
cond [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> Expr a -> [Text]
forall a. Expr a -> [Text]
getColumns Expr a
l [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> Expr a -> [Text]
forall a. Expr a -> [Text]
getColumns Expr a
r
getColumns (Unary op b a
_op Expr b
value) = Expr b -> [Text]
forall a. Expr a -> [Text]
getColumns Expr b
value
getColumns (Binary op c b a
_op Expr c
l Expr b
r) = Expr c -> [Text]
forall a. Expr a -> [Text]
getColumns Expr c
l [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> Expr b -> [Text]
forall a. Expr a -> [Text]
getColumns Expr b
r
getColumns (Agg AggStrategy a b
_strategy Expr b
expr) = Expr b -> [Text]
forall a. Expr a -> [Text]
getColumns Expr b
expr
getColumns (Over [Text]
keys Expr a
inner) = [Text]
keys [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> Expr a -> [Text]
forall a. Expr a -> [Text]
getColumns Expr a
inner

{- | Render an expression as readable, width-aware pseudo-code at the default
width ('P.defaultWidth'). See 'prettyPrintWidth' to control wrapping.
-}
prettyPrint :: Expr a -> String
prettyPrint :: forall a. Expr a -> String
prettyPrint = Int -> Expr a -> String
forall a. Int -> Expr a -> String
prettyPrintWidth Int
P.defaultWidth

{- | Render an expression as readable, width-aware pseudo-code: long binary chains
wrap onto aligned continuation lines, @if@/@then@/@else@ break onto their own lines
(nested @else if@ form a flat ladder), and sub-exprs are parenthesized by precedence.
-}
prettyPrintWidth :: Int -> Expr a -> String
prettyPrintWidth :: forall a. Int -> Expr a -> String
prettyPrintWidth Int
width = Int -> Doc -> String
P.render Int
width (Doc -> String) -> (Expr a -> Doc) -> Expr a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Expr a -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
0
  where
    toDoc :: Int -> Expr x -> P.Doc
    toDoc :: forall x. Int -> Expr x -> Doc
toDoc Int
prec Expr x
expr = case Expr x
expr of
        Col Text
name -> String -> Doc
P.text (Text -> String
T.unpack Text
name)
        CastWith Text
name Text
_ Either String a -> x
_ -> String -> Doc
P.text (Text -> String
T.unpack Text
name)
        CastExprWith Text
tag Either String a -> x
_ Expr src
inner -> String -> Doc
P.text (Text -> String
T.unpack Text
tag) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
P.parens (Int -> Expr src -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
0 Expr src
inner)
        Lit x
value -> String -> Doc
P.text (x -> String
forall a. Show a => a -> String
show x
value)
        If{} -> Int -> Expr x -> Doc
forall x. Int -> Expr x -> Doc
renderIf Int
prec Expr x
expr
        Unary op b x
op Expr b
arg ->
            let fn :: Text
fn = Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe (op b x -> Text
forall a b. op a b -> Text
forall (op :: * -> * -> *) a b. UnaryOp op => op a b -> Text
unaryName op b x
op) (op b x -> Maybe Text
forall a b. op a b -> Maybe Text
forall (op :: * -> * -> *) a b. UnaryOp op => op a b -> Maybe Text
unarySymbol op b x
op)
             in String -> Doc
P.text (Text -> String
T.unpack Text
fn) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
P.parens (Int -> Expr b -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
0 Expr b
arg)
        Binary op c b x
op Expr c
l Expr b
r -> case op c b x -> Maybe Text
forall a b c. op a b c -> Maybe Text
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> Maybe Text
binarySymbol op c b x
op of
            Just Text
sym -> Int -> op c b x -> String -> Expr c -> Expr b -> Doc
forall (op :: * -> * -> * -> *) c b a.
BinaryOp op =>
Int -> op c b a -> String -> Expr c -> Expr b -> Doc
renderBinary Int
prec op c b x
op (Text -> String
T.unpack Text
sym) Expr c
l Expr b
r
            Maybe Text
Nothing ->
                String -> Doc
P.text (Text -> String
T.unpack (op c b x -> Text
forall a b c. op a b c -> Text
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> Text
binaryName op c b x
op))
                    Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
P.parens (Int -> Expr c -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
0 Expr c
l Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
P.text String
", " Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Int -> Expr b -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
0 Expr b
r)
        Agg (CollectAgg Text
op v b -> x
_) Expr b
arg -> String -> Doc
P.text (Text -> String
T.unpack Text
op) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
P.parens (Int -> Expr b -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
0 Expr b
arg)
        Agg (FoldAgg Text
op Maybe x
_ x -> b -> x
_) Expr b
arg -> String -> Doc
P.text (Text -> String
T.unpack Text
op) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
P.parens (Int -> Expr b -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
0 Expr b
arg)
        Agg (MergeAgg Text
op acc
_ acc -> b -> acc
_ acc -> acc -> acc
_ acc -> x
_) Expr b
arg -> String -> Doc
P.text (Text -> String
T.unpack Text
op) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
P.parens (Int -> Expr b -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
0 Expr b
arg)
        Over [Text]
keys Expr x
inner ->
            Int -> Expr x -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
0 Expr x
inner Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
P.text (String
".over(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ [String] -> String
forall a. Show a => a -> String
show ((Text -> String) -> [Text] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Text -> String
T.unpack [Text]
keys) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")")

    renderBinary ::
        (BinaryOp op) => Int -> op c b a -> String -> Expr c -> Expr b -> P.Doc
    renderBinary :: forall (op :: * -> * -> * -> *) c b a.
BinaryOp op =>
Int -> op c b a -> String -> Expr c -> Expr b -> Doc
renderBinary Int
prec op c b a
op String
sym Expr c
l Expr b
r =
        let p :: Int
p = op c b a -> Int
forall a b c. op a b c -> Int
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> Int
binaryPrecedence op c b a
op
            body :: Doc
body
                | op c b a -> Bool
forall a b c. op a b c -> Bool
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> Bool
binaryCommutative op c b a
op =
                    let operands :: [Doc]
operands =
                            Text -> Int -> Expr c -> [Doc]
forall x. Text -> Int -> Expr x -> [Doc]
flattenChain (op c b a -> Text
forall a b c. op a b c -> Text
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> Text
binaryName op c b a
op) Int
p Expr c
l
                                [Doc] -> [Doc] -> [Doc]
forall a. [a] -> [a] -> [a]
++ Text -> Int -> Expr b -> [Doc]
forall x. Text -> Int -> Expr x -> [Doc]
flattenChain (op c b a -> Text
forall a b c. op a b c -> Text
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> Text
binaryName op c b a
op) Int
p Expr b
r
                     in case [Doc]
operands of
                            (Doc
o0 : os :: [Doc]
os@(Doc
_ : [Doc]
_)) ->
                                Doc -> Doc
P.group
                                    ( Doc
o0
                                        Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Int -> Doc -> Doc
P.nest
                                            Int
2
                                            ([Doc] -> Doc
P.hcat [Doc
P.line Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
P.text String
sym Doc -> Doc -> Doc
P.<+> Doc
o | Doc
o <- [Doc]
os])
                                    )
                            [Doc]
_ -> Int -> Expr c -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
p Expr c
l Doc -> Doc -> Doc
P.<+> String -> Doc
P.text String
sym Doc -> Doc -> Doc
P.<+> Int -> Expr b -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
p Expr b
r
                | Bool
otherwise =
                    Doc -> Doc
P.group (Int -> Expr c -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
p Expr c
l Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Int -> Doc -> Doc
P.nest Int
2 (Doc
P.line Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
P.text String
sym Doc -> Doc -> Doc
P.<+> Int -> Expr b -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
p Expr b
r))
         in if Int
prec Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
p
                then Doc -> Doc
P.parens Doc
body
                else if Int
prec Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1 Bool -> Bool -> Bool
&& Int
prec Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
p then Doc -> Doc
P.parensWhenBroken Doc
body else Doc
body

    flattenChain :: T.Text -> Int -> Expr x -> [P.Doc]
    flattenChain :: forall x. Text -> Int -> Expr x -> [Doc]
flattenChain Text
name Int
p Expr x
e = case Expr x
e of
        Binary op c b x
op' Expr c
l' Expr b
r'
            | op c b x -> Text
forall a b c. op a b c -> Text
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> Text
binaryName op c b x
op' Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
name
            , op c b x -> Int
forall a b c. op a b c -> Int
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> Int
binaryPrecedence op c b x
op' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
p
            , op c b x -> Bool
forall a b c. op a b c -> Bool
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> Bool
binaryCommutative op c b x
op' ->
                Text -> Int -> Expr c -> [Doc]
forall x. Text -> Int -> Expr x -> [Doc]
flattenChain Text
name Int
p Expr c
l' [Doc] -> [Doc] -> [Doc]
forall a. [a] -> [a] -> [a]
++ Text -> Int -> Expr b -> [Doc]
forall x. Text -> Int -> Expr x -> [Doc]
flattenChain Text
name Int
p Expr b
r'
        Expr x
_ -> [Int -> Expr x -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
p Expr x
e]

    renderIf :: Int -> Expr x -> P.Doc
    renderIf :: forall x. Int -> Expr x -> Doc
renderIf Int
prec (If Expr Bool
c Expr x
t Expr x
e) =
        let blk :: Doc
blk =
                String -> Doc
P.text String
"if" Doc -> Doc -> Doc
P.<+> Int -> Doc -> Doc
P.nest Int
3 (Doc -> Doc
P.group (Int -> Expr Bool -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
0 Expr Bool
c))
                    Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
P.hardline
                    Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
P.text String
"then" Doc -> Doc -> Doc
P.<+> Int -> Expr x -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
0 Expr x
t
                    Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
P.hardline
                    Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Expr x -> Doc
forall x. Expr x -> Doc
renderElse Expr x
e
         in if Int
prec Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 then Doc -> Doc
P.parens (Int -> Doc -> Doc
P.nest Int
2 Doc
blk) else Doc
blk
    renderIf Int
_ Expr x
_ = Doc
forall a. Monoid a => a
mempty

    renderElse :: Expr x -> P.Doc
    renderElse :: forall x. Expr x -> Doc
renderElse (If Expr Bool
c Expr x
t Expr x
e) =
        String -> Doc
P.text String
"else if" Doc -> Doc -> Doc
P.<+> Int -> Doc -> Doc
P.nest Int
8 (Doc -> Doc
P.group (Int -> Expr Bool -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
0 Expr Bool
c))
            Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
P.hardline
            Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
P.text String
"then" Doc -> Doc -> Doc
P.<+> Int -> Expr x -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
0 Expr x
t
            Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
P.hardline
            Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Expr x -> Doc
forall x. Expr x -> Doc
renderElse Expr x
e
    renderElse Expr x
other = String -> Doc
P.text String
"else" Doc -> Doc -> Doc
P.<+> Int -> Expr x -> Doc
forall x. Int -> Expr x -> Doc
toDoc Int
0 Expr x
other