{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}

module Jikka.RestrictedPython.Language.Builtin where

import Data.Functor
import qualified Data.Set as S
import Jikka.Common.Alpha
import Jikka.Common.Error
import Jikka.RestrictedPython.Language.Expr
import Jikka.RestrictedPython.Language.Util

builtinNames :: S.Set VarName
builtinNames :: Set VarName
builtinNames = Set VarName -> Set VarName -> Set VarName
forall a. Ord a => Set a -> Set a -> Set a
S.union Set VarName
standardBuiltinNames Set VarName
additionalBuiltinNames

standardBuiltinNames :: S.Set VarName
standardBuiltinNames :: Set VarName
standardBuiltinNames =
  [VarName] -> Set VarName
forall a. Ord a => [a] -> Set a
S.fromList
    [ VarName
"abs",
      VarName
"all",
      VarName
"any",
      VarName
"bool",
      VarName
"divmod",
      VarName
"enumerate",
      VarName
"filter",
      VarName
"input",
      VarName
"int",
      VarName
"len",
      VarName
"list",
      VarName
"map",
      VarName
"max",
      VarName
"min",
      VarName
"pow",
      VarName
"print",
      VarName
"range",
      VarName
"reversed",
      VarName
"sorted",
      VarName
"sum",
      VarName
"zip"
    ]

additionalBuiltinNames :: S.Set VarName
additionalBuiltinNames :: Set VarName
additionalBuiltinNames =
  [VarName] -> Set VarName
forall a. Ord a => [a] -> Set a
S.fromList
    [ VarName
"argmax",
      VarName
"argmin",
      VarName
"ceildiv",
      VarName
"ceilmod",
      VarName
"choose",
      VarName
"fact",
      VarName
"floordiv",
      VarName
"floormod",
      VarName
"gcd",
      VarName
"inv",
      VarName
"lcm",
      VarName
"multichoose",
      VarName
"permute",
      VarName
"product"
    ]

-- | `resolveUniqueBuiltin` makes a builtin function from a variable name.
-- However, this doesn't anything for ambiguous builtin functions.
-- For example, the builtin function "max" is kept as a variable because it may be \(\mathbf{list}(\alpha) \to \alpha\), \(\alpha \times \alpha \to \alpha\), etc. and this function cannot resolve it.
resolveUniqueBuiltin :: (MonadAlpha m, MonadError Error m) => VarName' -> m Expr'
resolveUniqueBuiltin :: VarName' -> m Expr'
resolveUniqueBuiltin VarName'
x | VarName' -> VarName
forall a. WithLoc' a -> a
value' VarName'
x VarName -> Set VarName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VarName
builtinNames = Expr' -> m Expr'
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr' -> m Expr') -> Expr' -> m Expr'
forall a b. (a -> b) -> a -> b
$ Maybe Loc -> Expr -> Expr'
forall a. Maybe Loc -> a -> WithLoc' a
WithLoc' (VarName' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' VarName'
x) (VarName' -> Expr
Name VarName'
x)
resolveUniqueBuiltin VarName'
x = do
  let f :: Builtin -> m Expr'
f = Expr' -> m Expr'
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr' -> m Expr') -> (Builtin -> Expr') -> Builtin -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe Loc -> Expr -> Expr'
forall a. Maybe Loc -> a -> WithLoc' a
WithLoc' (VarName' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' VarName'
x) (Expr -> Expr') -> (Builtin -> Expr) -> Builtin -> Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Constant -> Expr
Constant (Constant -> Expr) -> (Builtin -> Constant) -> Builtin -> Expr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builtin -> Constant
ConstBuiltin
  case VarName' -> VarName
forall a. WithLoc' a -> a
value' VarName'
x of
    VarName
"abs" -> Builtin -> m Expr'
f Builtin
BuiltinAbs
    VarName
"all" -> Builtin -> m Expr'
f Builtin
BuiltinAll
    VarName
"any" -> Builtin -> m Expr'
f Builtin
BuiltinAny
    VarName
"bool" -> Builtin -> m Expr'
f (Builtin -> m Expr') -> (Type -> Builtin) -> Type -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Builtin
BuiltinBool (Type -> m Expr') -> m Type -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
    VarName
"divmod" -> Builtin -> m Expr'
f Builtin
BuiltinDivMod
    VarName
"enumerate" -> Builtin -> m Expr'
f (Builtin -> m Expr') -> (Type -> Builtin) -> Type -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Builtin
BuiltinEnumerate (Type -> m Expr') -> m Type -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
    VarName
"filter" -> Builtin -> m Expr'
f (Builtin -> m Expr') -> (Type -> Builtin) -> Type -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Builtin
BuiltinFilter (Type -> m Expr') -> m Type -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
    VarName
"int" -> Builtin -> m Expr'
f (Builtin -> m Expr') -> (Type -> Builtin) -> Type -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Builtin
BuiltinInt (Type -> m Expr') -> m Type -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
    VarName
"input" -> Builtin -> m Expr'
f Builtin
BuiltinInput
    VarName
"len" -> Builtin -> m Expr'
f (Builtin -> m Expr') -> (Type -> Builtin) -> Type -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Builtin
BuiltinLen (Type -> m Expr') -> m Type -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
    VarName
"list" -> Builtin -> m Expr'
f (Builtin -> m Expr') -> (Type -> Builtin) -> Type -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Builtin
BuiltinList (Type -> m Expr') -> m Type -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
    VarName
"reversed" -> Builtin -> m Expr'
f (Builtin -> m Expr') -> (Type -> Builtin) -> Type -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Builtin
BuiltinReversed (Type -> m Expr') -> m Type -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
    VarName
"sorted" -> Builtin -> m Expr'
f (Builtin -> m Expr') -> (Type -> Builtin) -> Type -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Builtin
BuiltinSorted (Type -> m Expr') -> m Type -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
    VarName
"sum" -> Builtin -> m Expr'
f Builtin
BuiltinSum
    VarName
"argmax" -> Builtin -> m Expr'
f (Builtin -> m Expr') -> (Type -> Builtin) -> Type -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Builtin
BuiltinArgMax (Type -> m Expr') -> m Type -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
    VarName
"argmin" -> Builtin -> m Expr'
f (Builtin -> m Expr') -> (Type -> Builtin) -> Type -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Builtin
BuiltinArgMin (Type -> m Expr') -> m Type -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
    VarName
"ceildiv" -> Builtin -> m Expr'
f Builtin
BuiltinCeilDiv
    VarName
"ceilmod" -> Builtin -> m Expr'
f Builtin
BuiltinCeilMod
    VarName
"choose" -> Builtin -> m Expr'
f Builtin
BuiltinChoose
    VarName
"fact" -> Builtin -> m Expr'
f Builtin
BuiltinFact
    VarName
"floordiv" -> Builtin -> m Expr'
f Builtin
BuiltinFloorDiv
    VarName
"floormod" -> Builtin -> m Expr'
f Builtin
BuiltinFloorMod
    VarName
"gcd" -> Builtin -> m Expr'
f Builtin
BuiltinGcd
    VarName
"inv" -> Builtin -> m Expr'
f Builtin
BuiltinModInv
    VarName
"lcm" -> Builtin -> m Expr'
f Builtin
BuiltinLcm
    VarName
"multichoose" -> Builtin -> m Expr'
f Builtin
BuiltinMultiChoose
    VarName
"permute" -> Builtin -> m Expr'
f Builtin
BuiltinPermute
    VarName
"product" -> Builtin -> m Expr'
f Builtin
BuiltinProduct
    VarName
_ -> Expr' -> m Expr'
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr' -> m Expr') -> Expr' -> m Expr'
forall a b. (a -> b) -> a -> b
$ Maybe Loc -> Expr -> Expr'
forall a. Maybe Loc -> a -> WithLoc' a
WithLoc' (VarName' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' VarName'
x) (VarName' -> Expr
Name VarName'
x)

resolveBuiltin :: (MonadAlpha m, MonadError Error m) => VarName' -> Int -> m Expr'
resolveBuiltin :: VarName' -> Int -> m Expr'
resolveBuiltin VarName'
x Int
_ | VarName' -> VarName
forall a. WithLoc' a -> a
value' VarName'
x VarName -> Set VarName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VarName
builtinNames = Expr' -> m Expr'
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr' -> m Expr') -> Expr' -> m Expr'
forall a b. (a -> b) -> a -> b
$ Maybe Loc -> Expr -> Expr'
forall a. Maybe Loc -> a -> WithLoc' a
WithLoc' (VarName' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' VarName'
x) (VarName' -> Expr
Name VarName'
x)
resolveBuiltin VarName'
x Int
n = Maybe Loc -> m Expr' -> m Expr'
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> m a -> m a
wrapAt' (VarName' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' VarName'
x) (m Expr' -> m Expr') -> (m Expr' -> m Expr') -> m Expr' -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> m Expr' -> m Expr'
forall (m :: * -> *) a. MonadError Error m => String -> m a -> m a
wrapError' String
"Jikka.RestrictedPython.Language.Builtin.resolveBuiltin" (m Expr' -> m Expr') -> m Expr' -> m Expr'
forall a b. (a -> b) -> a -> b
$ do
  let f :: Builtin -> m Expr'
f = Expr' -> m Expr'
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr' -> m Expr') -> (Builtin -> Expr') -> Builtin -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe Loc -> Expr -> Expr'
forall a. Maybe Loc -> a -> WithLoc' a
WithLoc' (VarName' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' VarName'
x) (Expr -> Expr') -> (Builtin -> Expr) -> Builtin -> Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Constant -> Expr
Constant (Constant -> Expr) -> (Builtin -> Constant) -> Builtin -> Expr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builtin -> Constant
ConstBuiltin
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    String -> m ()
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwInternalError (String -> m ()) -> String -> m ()
forall a b. (a -> b) -> a -> b
$ String
"negative arity: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n
  case VarName' -> VarName
forall a. WithLoc' a -> a
value' VarName'
x of
    VarName
"map" -> Builtin -> m Expr'
f (Builtin -> m Expr') -> m Builtin -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ([Type] -> Type -> Builtin
BuiltinMap ([Type] -> Type -> Builtin) -> m [Type] -> m (Type -> Builtin)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m Type -> m [Type]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType m (Type -> Builtin) -> m Type -> m Builtin
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType)
    VarName
"max" -> case Int
n of
      Int
1 -> Builtin -> m Expr'
f (Builtin -> m Expr') -> (Type -> Builtin) -> Type -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Builtin
BuiltinMax1 (Type -> m Expr') -> m Type -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
      Int
_ -> Builtin -> m Expr'
f (Builtin -> m Expr') -> m Builtin -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Type -> Int -> Builtin
BuiltinMax (Type -> Int -> Builtin) -> m Type -> m (Int -> Builtin)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType m (Int -> Builtin) -> m Int -> m Builtin
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> m Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
n)
    VarName
"min" -> case Int
n of
      Int
1 -> Builtin -> m Expr'
f (Builtin -> m Expr') -> (Type -> Builtin) -> Type -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Builtin
BuiltinMin1 (Type -> m Expr') -> m Type -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
      Int
_ -> Builtin -> m Expr'
f (Builtin -> m Expr') -> m Builtin -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Type -> Int -> Builtin
BuiltinMin (Type -> Int -> Builtin) -> m Type -> m (Int -> Builtin)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType m (Int -> Builtin) -> m Int -> m Builtin
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> m Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
n)
    VarName
"pow" ->
      if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
3
        then Builtin -> m Expr'
f Builtin
BuiltinModPow
        else Builtin -> m Expr'
f Builtin
BuiltinPow
    VarName
"print" -> Builtin -> m Expr'
f (Builtin -> m Expr') -> ([Type] -> Builtin) -> [Type] -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Type] -> Builtin
BuiltinPrint ([Type] -> m Expr') -> m [Type] -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> m Type -> m [Type]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
    VarName
"range" -> case Int
n of
      Int
1 -> Builtin -> m Expr'
f Builtin
BuiltinRange1
      Int
2 -> Builtin -> m Expr'
f Builtin
BuiltinRange2
      Int
3 -> Builtin -> m Expr'
f Builtin
BuiltinRange3
      Int
_ -> String -> m Expr'
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwTypeError (String -> m Expr') -> String -> m Expr'
forall a b. (a -> b) -> a -> b
$ String
"range expected 1, 2, or 3 arguments, got " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n
    VarName
"zip" -> Builtin -> m Expr'
f (Builtin -> m Expr') -> ([Type] -> Builtin) -> [Type] -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Type] -> Builtin
BuiltinZip ([Type] -> m Expr') -> m [Type] -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> m Type -> m [Type]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
    VarName
_ -> do
      Expr'
e <- VarName' -> m Expr'
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
VarName' -> m Expr'
resolveUniqueBuiltin VarName'
x
      case Expr' -> Expr
forall a. WithLoc' a -> a
value' Expr'
e of
        Constant (ConstBuiltin Builtin
_) -> Expr' -> m Expr'
forall (m :: * -> *) a. Monad m => a -> m a
return Expr'
e
        Expr
_ -> String -> m Expr'
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwInternalError (String -> m Expr') -> String -> m Expr'
forall a b. (a -> b) -> a -> b
$ String
"not exhaustive: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ VarName -> String
unVarName (VarName' -> VarName
forall a. WithLoc' a -> a
value' VarName'
x)

formatBuiltin :: Builtin -> String
formatBuiltin :: Builtin -> String
formatBuiltin = \case
  Builtin
BuiltinAbs -> String
"abs"
  Builtin
BuiltinPow -> String
"pow"
  Builtin
BuiltinModPow -> String
"pow"
  Builtin
BuiltinAll -> String
"all"
  Builtin
BuiltinAny -> String
"any"
  Builtin
BuiltinDivMod -> String
"divmod"
  BuiltinSorted Type
_ -> String
"sorted"
  BuiltinEnumerate Type
_ -> String
"enumerate"
  BuiltinBool Type
_ -> String
"bool"
  BuiltinInt Type
_ -> String
"int"
  Builtin
BuiltinSum -> String
"sum"
  BuiltinZip [Type]
_ -> String
"zip"
  BuiltinFilter Type
_ -> String
"filter"
  BuiltinTuple [Type]
_ -> String
"tuple"
  BuiltinLen Type
_ -> String
"len"
  BuiltinList Type
_ -> String
"list"
  Builtin
BuiltinRange1 -> String
"range"
  Builtin
BuiltinRange2 -> String
"range"
  Builtin
BuiltinRange3 -> String
"range"
  BuiltinMap [Type]
_ Type
_ -> String
"map"
  BuiltinReversed Type
_ -> String
"reversed"
  BuiltinMax1 Type
_ -> String
"max"
  BuiltinMax Type
_ Int
_ -> String
"max"
  BuiltinMin1 Type
_ -> String
"min"
  BuiltinMin Type
_ Int
_ -> String
"min"
  BuiltinArgMax Type
_ -> String
"argmax"
  BuiltinArgMin Type
_ -> String
"argmin"
  Builtin
BuiltinCeilDiv -> String
"ceildiv"
  Builtin
BuiltinCeilMod -> String
"ceilmod"
  Builtin
BuiltinFloorDiv -> String
"floordiv"
  Builtin
BuiltinFloorMod -> String
"floormod"
  Builtin
BuiltinChoose -> String
"choose"
  Builtin
BuiltinFact -> String
"fact"
  Builtin
BuiltinGcd -> String
"gcd"
  Builtin
BuiltinLcm -> String
"lcm"
  Builtin
BuiltinModInv -> String
"inv"
  Builtin
BuiltinMultiChoose -> String
"multichoose"
  Builtin
BuiltinPermute -> String
"permute"
  Builtin
BuiltinProduct -> String
"product"
  Builtin
BuiltinInput -> String
"input"
  BuiltinPrint [Type]
_ -> String
"print"

typeBuiltin :: Builtin -> Type
typeBuiltin :: Builtin -> Type
typeBuiltin = \case
  Builtin
BuiltinAbs -> [Type] -> Type -> Type
CallableTy [Type
IntTy] Type
IntTy
  Builtin
BuiltinPow -> [Type] -> Type -> Type
CallableTy [Type
IntTy] Type
IntTy
  Builtin
BuiltinModPow -> [Type] -> Type -> Type
CallableTy [Type
IntTy, Type
IntTy] Type
IntTy
  Builtin
BuiltinAll -> [Type] -> Type -> Type
CallableTy [Type -> Type
ListTy Type
BoolTy] Type
BoolTy
  Builtin
BuiltinAny -> [Type] -> Type -> Type
CallableTy [Type -> Type
ListTy Type
BoolTy] Type
BoolTy
  BuiltinArgMax Type
t -> [Type] -> Type -> Type
CallableTy [Type -> Type
ListTy Type
t] Type
IntTy
  BuiltinArgMin Type
t -> [Type] -> Type -> Type
CallableTy [Type -> Type
ListTy Type
t] Type
IntTy
  BuiltinBool Type
t -> [Type] -> Type -> Type
CallableTy [Type
t] Type
BoolTy
  Builtin
BuiltinCeilDiv -> [Type] -> Type -> Type
CallableTy [Type
IntTy, Type
IntTy] Type
IntTy
  Builtin
BuiltinCeilMod -> [Type] -> Type -> Type
CallableTy [Type
IntTy, Type
IntTy] Type
IntTy
  Builtin
BuiltinChoose -> [Type] -> Type -> Type
CallableTy [Type
IntTy, Type
IntTy] Type
IntTy
  Builtin
BuiltinDivMod -> [Type] -> Type -> Type
CallableTy [Type
IntTy, Type
IntTy] ([Type] -> Type
TupleTy [Type
IntTy, Type
IntTy])
  BuiltinEnumerate Type
t -> [Type] -> Type -> Type
CallableTy [Type -> Type
ListTy Type
t] (Type -> Type
ListTy ([Type] -> Type
TupleTy [Type
IntTy, Type
t]))
  Builtin
BuiltinFact -> [Type] -> Type -> Type
CallableTy [Type -> Type
ListTy Type
IntTy] Type
IntTy
  BuiltinFilter Type
t -> [Type] -> Type -> Type
CallableTy [[Type] -> Type -> Type
CallableTy [Type
t] Type
BoolTy, Type -> Type
ListTy Type
t] (Type -> Type
ListTy Type
t)
  Builtin
BuiltinFloorDiv -> [Type] -> Type -> Type
CallableTy [Type
IntTy, Type
IntTy] Type
IntTy
  Builtin
BuiltinFloorMod -> [Type] -> Type -> Type
CallableTy [Type
IntTy, Type
IntTy] Type
IntTy
  Builtin
BuiltinGcd -> [Type] -> Type -> Type
CallableTy [Type
IntTy, Type
IntTy] Type
IntTy
  BuiltinInt Type
t -> [Type] -> Type -> Type
CallableTy [Type
t] Type
IntTy
  Builtin
BuiltinModInv -> [Type] -> Type -> Type
CallableTy [Type
IntTy, Type
IntTy] Type
IntTy
  Builtin
BuiltinLcm -> [Type] -> Type -> Type
CallableTy [Type
IntTy, Type
IntTy] Type
IntTy
  BuiltinLen Type
t -> [Type] -> Type -> Type
CallableTy [Type -> Type
ListTy Type
t] Type
IntTy
  BuiltinList Type
t -> [Type] -> Type -> Type
CallableTy [Type -> Type
ListTy Type
t] (Type -> Type
ListTy Type
t)
  BuiltinMap [Type]
args Type
ret -> [Type] -> Type -> Type
CallableTy ([Type] -> Type -> Type
CallableTy [Type]
args Type
ret Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
ListTy [Type]
args) (Type -> Type
ListTy Type
ret)
  BuiltinMax Type
t Int
n -> [Type] -> Type -> Type
CallableTy (Int -> Type -> [Type]
forall a. Int -> a -> [a]
replicate Int
n Type
t) Type
t
  BuiltinMax1 Type
t -> [Type] -> Type -> Type
CallableTy [Type -> Type
ListTy Type
t] Type
t
  BuiltinMin Type
t Int
n -> [Type] -> Type -> Type
CallableTy (Int -> Type -> [Type]
forall a. Int -> a -> [a]
replicate Int
n Type
t) Type
t
  BuiltinMin1 Type
t -> [Type] -> Type -> Type
CallableTy [Type -> Type
ListTy Type
t] Type
t
  Builtin
BuiltinMultiChoose -> [Type] -> Type -> Type
CallableTy [Type
IntTy, Type
IntTy] Type
IntTy
  Builtin
BuiltinPermute -> [Type] -> Type -> Type
CallableTy [Type
IntTy, Type
IntTy] Type
IntTy
  Builtin
BuiltinProduct -> [Type] -> Type -> Type
CallableTy [Type -> Type
ListTy Type
IntTy] Type
IntTy
  Builtin
BuiltinRange1 -> [Type] -> Type -> Type
CallableTy [Type
IntTy] (Type -> Type
ListTy Type
IntTy)
  Builtin
BuiltinRange2 -> [Type] -> Type -> Type
CallableTy [Type
IntTy, Type
IntTy] (Type -> Type
ListTy Type
IntTy)
  Builtin
BuiltinRange3 -> [Type] -> Type -> Type
CallableTy [Type
IntTy, Type
IntTy, Type
IntTy] (Type -> Type
ListTy Type
IntTy)
  BuiltinReversed Type
t -> [Type] -> Type -> Type
CallableTy [Type -> Type
ListTy Type
t] (Type -> Type
ListTy Type
t)
  BuiltinSorted Type
t -> [Type] -> Type -> Type
CallableTy [Type -> Type
ListTy Type
t] (Type -> Type
ListTy Type
t)
  Builtin
BuiltinSum -> [Type] -> Type -> Type
CallableTy [Type -> Type
ListTy Type
IntTy] Type
IntTy
  BuiltinTuple [Type]
ts -> [Type] -> Type -> Type
CallableTy [[Type] -> Type
TupleTy [Type]
ts] ([Type] -> Type
TupleTy [Type]
ts)
  BuiltinZip [Type]
ts -> [Type] -> Type -> Type
CallableTy ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
ListTy [Type]
ts) ([Type] -> Type
TupleTy [Type]
ts)
  Builtin
BuiltinInput -> [Type] -> Type -> Type
CallableTy [] Type
StringTy
  BuiltinPrint [Type]
ts -> [Type] -> Type -> Type
CallableTy [Type]
ts Type
SideEffectTy

mapTypeBuiltin :: (Type -> Type) -> Builtin -> Builtin
mapTypeBuiltin :: (Type -> Type) -> Builtin -> Builtin
mapTypeBuiltin Type -> Type
f = \case
  Builtin
BuiltinAbs -> Builtin
BuiltinAbs
  Builtin
BuiltinPow -> Builtin
BuiltinPow
  Builtin
BuiltinModPow -> Builtin
BuiltinModPow
  Builtin
BuiltinAll -> Builtin
BuiltinAll
  Builtin
BuiltinAny -> Builtin
BuiltinAny
  BuiltinArgMax Type
t -> Type -> Builtin
BuiltinArgMax (Type -> Type
f Type
t)
  BuiltinArgMin Type
t -> Type -> Builtin
BuiltinArgMin (Type -> Type
f Type
t)
  BuiltinBool Type
t -> Type -> Builtin
BuiltinBool (Type -> Type
f Type
t)
  Builtin
BuiltinCeilDiv -> Builtin
BuiltinCeilDiv
  Builtin
BuiltinCeilMod -> Builtin
BuiltinCeilMod
  Builtin
BuiltinChoose -> Builtin
BuiltinChoose
  Builtin
BuiltinDivMod -> Builtin
BuiltinDivMod
  BuiltinEnumerate Type
t -> Type -> Builtin
BuiltinEnumerate (Type -> Type
f Type
t)
  Builtin
BuiltinFact -> Builtin
BuiltinFact
  BuiltinFilter Type
t -> Type -> Builtin
BuiltinFilter (Type -> Type
f Type
t)
  Builtin
BuiltinFloorDiv -> Builtin
BuiltinFloorDiv
  Builtin
BuiltinFloorMod -> Builtin
BuiltinFloorMod
  Builtin
BuiltinGcd -> Builtin
BuiltinGcd
  BuiltinInt Type
t -> Type -> Builtin
BuiltinInt (Type -> Type
f Type
t)
  Builtin
BuiltinModInv -> Builtin
BuiltinModInv
  Builtin
BuiltinLcm -> Builtin
BuiltinLcm
  BuiltinLen Type
t -> Type -> Builtin
BuiltinLen (Type -> Type
f Type
t)
  BuiltinList Type
t -> Type -> Builtin
BuiltinList (Type -> Type
f Type
t)
  BuiltinMap [Type]
args Type
ret -> [Type] -> Type -> Builtin
BuiltinMap ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
f [Type]
args) (Type -> Type
f Type
ret)
  BuiltinMax Type
t Int
n -> Type -> Int -> Builtin
BuiltinMax (Type -> Type
f Type
t) Int
n
  BuiltinMax1 Type
t -> Type -> Builtin
BuiltinMax1 (Type -> Type
f Type
t)
  BuiltinMin Type
t Int
n -> Type -> Int -> Builtin
BuiltinMin (Type -> Type
f Type
t) Int
n
  BuiltinMin1 Type
t -> Type -> Builtin
BuiltinMin1 (Type -> Type
f Type
t)
  Builtin
BuiltinMultiChoose -> Builtin
BuiltinMultiChoose
  Builtin
BuiltinPermute -> Builtin
BuiltinPermute
  Builtin
BuiltinProduct -> Builtin
BuiltinProduct
  Builtin
BuiltinRange1 -> Builtin
BuiltinRange1
  Builtin
BuiltinRange2 -> Builtin
BuiltinRange2
  Builtin
BuiltinRange3 -> Builtin
BuiltinRange3
  BuiltinReversed Type
t -> Type -> Builtin
BuiltinReversed (Type -> Type
f Type
t)
  BuiltinSorted Type
t -> Type -> Builtin
BuiltinSorted (Type -> Type
f Type
t)
  Builtin
BuiltinSum -> Builtin
BuiltinSum
  BuiltinTuple [Type]
ts -> [Type] -> Builtin
BuiltinTuple ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
f [Type]
ts)
  BuiltinZip [Type]
ts -> [Type] -> Builtin
BuiltinZip ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
f [Type]
ts)
  Builtin
BuiltinInput -> Builtin
BuiltinInput
  BuiltinPrint [Type]
ts -> [Type] -> Builtin
BuiltinPrint ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
f [Type]
ts)

attributeNames :: S.Set AttributeName
attributeNames :: Set AttributeName
attributeNames =
  [AttributeName] -> Set AttributeName
forall a. Ord a => [a] -> Set a
S.fromList
    [ AttributeName
"count",
      AttributeName
"index",
      AttributeName
"copy",
      AttributeName
"append",
      AttributeName
"split"
    ]

resolveAttribute' :: (MonadAlpha m, MonadError Error m) => Attribute' -> m Attribute'
resolveAttribute' :: Attribute' -> m Attribute'
resolveAttribute' Attribute'
x = Maybe Loc -> m Attribute' -> m Attribute'
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> m a -> m a
wrapAt' (Attribute' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Attribute'
x) (m Attribute' -> m Attribute') -> m Attribute' -> m Attribute'
forall a b. (a -> b) -> a -> b
$ case Attribute' -> Attribute
forall a. WithLoc' a -> a
value' Attribute'
x of
  UnresolvedAttribute AttributeName
x' ->
    if AttributeName
x' AttributeName -> Set AttributeName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set AttributeName
attributeNames
      then String -> m Attribute'
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwSymbolError (String -> m Attribute') -> String -> m Attribute'
forall a b. (a -> b) -> a -> b
$ String
"unknown attribute: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ AttributeName -> String
unAttributeName AttributeName
x'
      else String -> m Attribute' -> m Attribute'
forall (m :: * -> *) a. MonadError Error m => String -> m a -> m a
wrapError' String
"Jikka.RestrictedPython.Language.Builtin.resolveAttribute" (m Attribute' -> m Attribute') -> m Attribute' -> m Attribute'
forall a b. (a -> b) -> a -> b
$ do
        Maybe Loc -> Attribute -> Attribute'
forall a. Maybe Loc -> a -> WithLoc' a
WithLoc' (Attribute' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Attribute'
x) (Attribute -> Attribute') -> m Attribute -> m Attribute'
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case AttributeName
x' of
          AttributeName
"count" -> Type -> Attribute
BuiltinCount (Type -> Attribute) -> m Type -> m Attribute
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
          AttributeName
"index" -> Type -> Attribute
BuiltinIndex (Type -> Attribute) -> m Type -> m Attribute
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
          AttributeName
"copy" -> Type -> Attribute
BuiltinCopy (Type -> Attribute) -> m Type -> m Attribute
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
          AttributeName
"append" -> Type -> Attribute
BuiltinAppend (Type -> Attribute) -> m Type -> m Attribute
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Type
forall (m :: * -> *). MonadAlpha m => m Type
genType
          AttributeName
"split" -> Attribute -> m Attribute
forall (m :: * -> *) a. Monad m => a -> m a
return Attribute
BuiltinSplit
          AttributeName
_ -> String -> m Attribute
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwInternalError (String -> m Attribute) -> String -> m Attribute
forall a b. (a -> b) -> a -> b
$ String
"not exhaustive: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ AttributeName -> String
unAttributeName AttributeName
x'
  Attribute
_ -> Attribute' -> m Attribute'
forall (m :: * -> *) a. Monad m => a -> m a
return Attribute'
x

resolveAttribute :: (MonadAlpha m, MonadError Error m) => Expr' -> Attribute' -> m Expr
resolveAttribute :: Expr' -> Attribute' -> m Expr
resolveAttribute e :: Expr'
e@(WithLoc' Maybe Loc
_ (Name (WithLoc' Maybe Loc
_ VarName
"math"))) Attribute'
x = Maybe Loc -> m Expr -> m Expr
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> m a -> m a
wrapAt' (Attribute' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Attribute'
x) (m Expr -> m Expr) -> m Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ case Attribute' -> Attribute
forall a. WithLoc' a -> a
value' Attribute'
x of
  UnresolvedAttribute AttributeName
x' -> case AttributeName
x' of
    AttributeName
"gcd" -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Constant -> Expr
Constant (Builtin -> Constant
ConstBuiltin Builtin
BuiltinGcd))
    AttributeName
"lcm" -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Constant -> Expr
Constant (Builtin -> Constant
ConstBuiltin Builtin
BuiltinGcd))
    AttributeName
_ -> String -> m Expr
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwSymbolError (String -> m Expr) -> String -> m Expr
forall a b. (a -> b) -> a -> b
$ String
"unknown attribute: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ AttributeName -> String
unAttributeName AttributeName
x'
  Attribute
_ -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ Expr' -> Attribute' -> Expr
Attribute Expr'
e Attribute'
x
resolveAttribute e :: Expr'
e@(WithLoc' Maybe Loc
_ (Name (WithLoc' Maybe Loc
_ VarName
"jikka"))) Attribute'
x = Maybe Loc -> m Expr -> m Expr
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> m a -> m a
wrapAt' (Attribute' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Attribute'
x) (m Expr -> m Expr) -> m Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ case Attribute' -> Attribute
forall a. WithLoc' a -> a
value' Attribute'
x of
  UnresolvedAttribute AttributeName
x' ->
    let x'' :: VarName
x'' = String -> VarName
VarName (AttributeName -> String
unAttributeName AttributeName
x')
     in if VarName
x'' VarName -> Set VarName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VarName
additionalBuiltinNames
          then String -> m Expr
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwSymbolError (String -> m Expr) -> String -> m Expr
forall a b. (a -> b) -> a -> b
$ String
"unknown attribute: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ AttributeName -> String
unAttributeName AttributeName
x'
          else Expr' -> Expr
forall a. WithLoc' a -> a
value' (Expr' -> Expr) -> m Expr' -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VarName' -> m Expr'
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
VarName' -> m Expr'
resolveUniqueBuiltin (Attribute'
x Attribute' -> VarName -> VarName'
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> VarName
x'')
  Attribute
_ -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ Expr' -> Attribute' -> Expr
Attribute Expr'
e Attribute'
x
resolveAttribute Expr'
e Attribute'
x = Expr' -> Attribute' -> Expr
Attribute Expr'
e (Attribute' -> Expr) -> m Attribute' -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Attribute' -> m Attribute'
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
Attribute' -> m Attribute'
resolveAttribute' Attribute'
x

formatAttribute :: Attribute -> String
formatAttribute :: Attribute -> String
formatAttribute = \case
  UnresolvedAttribute AttributeName
x -> AttributeName -> String
unAttributeName AttributeName
x
  BuiltinCount Type
_ -> String
"count"
  BuiltinIndex Type
_ -> String
"index"
  BuiltinCopy Type
_ -> String
"copy"
  BuiltinAppend Type
_ -> String
"append"
  Attribute
BuiltinSplit -> String
"split"

typeAttribute :: Attribute -> (Type, Type)
typeAttribute :: Attribute -> (Type, Type)
typeAttribute = \case
  UnresolvedAttribute AttributeName
x -> String -> (Type, Type)
forall a. HasCallStack => String -> a
error (String -> (Type, Type)) -> String -> (Type, Type)
forall a b. (a -> b) -> a -> b
$ String
"Jikka.RestrictedPython.Language.Builtin.typeAttribute: attributes must be resolved: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ AttributeName -> String
unAttributeName AttributeName
x
  BuiltinCount Type
t -> (Type -> Type
ListTy Type
t, [Type] -> Type -> Type
CallableTy [Type
t] Type
IntTy)
  BuiltinIndex Type
t -> (Type -> Type
ListTy Type
t, [Type] -> Type -> Type
CallableTy [Type
t] Type
IntTy)
  BuiltinCopy Type
t -> (Type -> Type
ListTy Type
t, [Type] -> Type -> Type
CallableTy [] (Type -> Type
ListTy Type
t))
  BuiltinAppend Type
t -> (Type -> Type
ListTy Type
t, [Type] -> Type -> Type
CallableTy [Type
t] Type
SideEffectTy)
  Attribute
BuiltinSplit -> (Type
StringTy, [Type] -> Type -> Type
CallableTy [] (Type -> Type
ListTy Type
StringTy))

mapTypeAttribute :: (Type -> Type) -> Attribute -> Attribute
mapTypeAttribute :: (Type -> Type) -> Attribute -> Attribute
mapTypeAttribute Type -> Type
f = \case
  UnresolvedAttribute AttributeName
x -> AttributeName -> Attribute
UnresolvedAttribute AttributeName
x
  BuiltinCount Type
t -> Type -> Attribute
BuiltinCount (Type -> Type
f Type
t)
  BuiltinIndex Type
t -> Type -> Attribute
BuiltinIndex (Type -> Type
f Type
t)
  BuiltinCopy Type
t -> Type -> Attribute
BuiltinCopy (Type -> Type
f Type
t)
  BuiltinAppend Type
t -> Type -> Attribute
BuiltinAppend (Type -> Type
f Type
t)
  Attribute
BuiltinSplit -> Attribute
BuiltinSplit