{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}

-- |
-- Module      : Jikka.RestrictedPython.Convert.ParseMain
-- Description : analyze @main@ function into input formats. / @main@ 関数を分析して入力フォーマットを得ます。
-- Copyright   : (c) Kimiyuki Onaka, 2021
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
module Jikka.RestrictedPython.Convert.ParseMain
  ( run,
  )
where

import Control.Arrow
import Data.Maybe
import Jikka.Common.Alpha
import Jikka.Common.Error
import Jikka.Common.IOFormat
import Jikka.RestrictedPython.Format (formatExpr, formatTarget)
import Jikka.RestrictedPython.Language.Expr
import Jikka.RestrictedPython.Language.Util

type MainFunction = (Maybe Loc, [(VarName', Type)], Type, [Statement])

splitMain :: Program -> (Maybe MainFunction, Program)
splitMain :: Program -> (Maybe MainFunction, Program)
splitMain = \case
  [] -> (Maybe MainFunction
forall a. Maybe a
Nothing, [])
  ToplevelFunctionDef (WithLoc' Maybe Loc
loc (VarName [Char]
"main")) [(WithLoc' VarName, Type)]
args Type
ret [Statement]
body : Program
stmts -> (MainFunction -> Maybe MainFunction
forall a. a -> Maybe a
Just (Maybe Loc
loc, [(WithLoc' VarName, Type)]
args, Type
ret, [Statement]
body), Program
stmts)
  ToplevelStatement
stmt : Program
stmts -> (Program -> Program)
-> (Maybe MainFunction, Program) -> (Maybe MainFunction, Program)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (ToplevelStatement
stmt ToplevelStatement -> Program -> Program
forall a. a -> [a] -> [a]
:) ((Maybe MainFunction, Program) -> (Maybe MainFunction, Program))
-> (Maybe MainFunction, Program) -> (Maybe MainFunction, Program)
forall a b. (a -> b) -> a -> b
$ Program -> (Maybe MainFunction, Program)
splitMain Program
stmts

checkMainType :: MonadError Error m => MainFunction -> m ()
checkMainType :: MainFunction -> m ()
checkMainType (Maybe Loc
loc, [(WithLoc' VarName, Type)]
args, Type
ret, [Statement]
_) = Maybe Loc -> m () -> m ()
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> m a -> m a
wrapAt' Maybe Loc
loc (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ case [(WithLoc' VarName, Type)]
args of
  (WithLoc' VarName, Type)
_ : [(WithLoc' VarName, Type)]
_ -> [Char] -> m ()
forall (m :: * -> *) a. MonadError Error m => [Char] -> m a
throwTypeError [Char]
"main function must not take arguments"
  [] -> case Type
ret of
    VarTy TypeName
_ -> () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Type
NoneTy -> () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Type
_ -> [Char] -> m ()
forall (m :: * -> *) a. MonadError Error m => [Char] -> m a
throwTypeError [Char]
"main function must return None"

pattern $mCallBuiltin :: forall r.
WithLoc' Expr
-> (Builtin -> [WithLoc' Expr] -> r) -> (Void# -> r) -> r
CallBuiltin b args <- WithLoc' _ (Call (WithLoc' _ (Constant (ConstBuiltin b))) args)

pattern $mCallMethod :: forall r.
WithLoc' Expr
-> (WithLoc' Expr -> Attribute' -> [WithLoc' Expr] -> r)
-> (Void# -> r)
-> r
CallMethod e a args <- WithLoc' _ (Call (WithLoc' _ (Attribute e a)) args)

pattern $mIntInput :: forall r. WithLoc' Expr -> (Void# -> r) -> (Void# -> r) -> r
IntInput <-
  CallBuiltin (BuiltinInt _) [CallBuiltin BuiltinInput []]

pattern $mMapIntInputSplit :: forall r. WithLoc' Expr -> (Void# -> r) -> (Void# -> r) -> r
MapIntInputSplit <-
  CallBuiltin
    (BuiltinMap [_] _)
    [ WithLoc' _ (Constant (ConstBuiltin (BuiltinInt _))),
      CallMethod
        (CallBuiltin BuiltinInput [])
        (WithLoc' _ BuiltinSplit)
        []
      ]

pattern $mListMapIntInputSplit :: forall r. WithLoc' Expr -> (Void# -> r) -> (Void# -> r) -> r
ListMapIntInputSplit <-
  CallBuiltin
    (BuiltinList _)
    [ CallBuiltin
        (BuiltinMap [_] _)
        [ WithLoc' _ (Constant (ConstBuiltin (BuiltinInt _))),
          CallMethod
            (CallBuiltin BuiltinInput [])
            (WithLoc' _ BuiltinSplit)
            []
          ]
      ]

pattern $mListRange :: forall r. WithLoc' Expr -> (VarName -> r) -> (Void# -> r) -> r
ListRange n <-
  CallBuiltin
    (BuiltinList _)
    [CallBuiltin BuiltinRange1 [WithLoc' _ (Name (WithLoc' _ n))]]

parseAnnAssign :: (MonadAlpha m, MonadError Error m) => Target' -> Type -> Expr' -> [Statement] -> m (FormatTree, Maybe ([String], Either String [String]), [Statement])
parseAnnAssign :: Target'
-> Type
-> WithLoc' Expr
-> [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      [Statement])
parseAnnAssign Target'
x Type
_ WithLoc' Expr
e [Statement]
cont = do
  let subscriptTrg :: Target' -> m ([Char], [[Char]])
subscriptTrg Target'
x = case Target' -> Target
forall a. WithLoc' a -> a
value' Target'
x of
        NameTrg WithLoc' VarName
x -> ([Char], [[Char]]) -> m ([Char], [[Char]])
forall (m :: * -> *) a. Monad m => a -> m a
return (VarName -> [Char]
unVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
x), [])
        SubscriptTrg Target'
x (WithLoc' Maybe Loc
_ (Name WithLoc' VarName
i)) -> ([[Char]] -> [[Char]]) -> ([Char], [[Char]]) -> ([Char], [[Char]])
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ([[Char]] -> [[Char]] -> [[Char]]
forall a. [a] -> [a] -> [a]
++ [VarName -> [Char]
unVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
i)]) (([Char], [[Char]]) -> ([Char], [[Char]]))
-> m ([Char], [[Char]]) -> m ([Char], [[Char]])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Target' -> m ([Char], [[Char]])
subscriptTrg Target'
x
        Target
_ -> Maybe Loc -> [Char] -> m ([Char], [[Char]])
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> [Char] -> m a
throwSemanticErrorAt' (Target' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Target'
x) ([Char] -> m ([Char], [[Char]])) -> [Char] -> m ([Char], [[Char]])
forall a b. (a -> b) -> a -> b
$ [Char]
"name target or subscript target is expected, but got: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Target' -> [Char]
formatTarget Target'
x
  let subscriptTupleTrg :: Target' -> m [([Char], [[Char]])]
subscriptTupleTrg Target'
x = case Target' -> Target
forall a. WithLoc' a -> a
value' Target'
x of
        TupleTrg [Target']
xs -> (Target' -> m ([Char], [[Char]]))
-> [Target'] -> m [([Char], [[Char]])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Target' -> m ([Char], [[Char]])
forall (m :: * -> *).
MonadError Error m =>
Target' -> m ([Char], [[Char]])
subscriptTrg [Target']
xs
        Target
_ -> Maybe Loc -> [Char] -> m [([Char], [[Char]])]
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> [Char] -> m a
throwSemanticErrorAt' (Target' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Target'
x) ([Char] -> m [([Char], [[Char]])])
-> [Char] -> m [([Char], [[Char]])]
forall a b. (a -> b) -> a -> b
$ [Char]
"tuple target is expected, but got: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Target' -> [Char]
formatTarget Target'
x
  let nameTrg :: Target' -> m [Char]
nameTrg Target'
x = case Target' -> Target
forall a. WithLoc' a -> a
value' Target'
x of
        NameTrg WithLoc' VarName
x -> [Char] -> m [Char]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Char] -> m [Char]) -> [Char] -> m [Char]
forall a b. (a -> b) -> a -> b
$ VarName -> [Char]
unVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
x)
        Target
_ -> Maybe Loc -> [Char] -> m [Char]
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> [Char] -> m a
throwSemanticErrorAt' (Target' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Target'
x) ([Char] -> m [Char]) -> [Char] -> m [Char]
forall a b. (a -> b) -> a -> b
$ [Char]
"name target is expected, but got: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Target' -> [Char]
formatTarget Target'
x
  let nameOrTupleTrg :: Target' -> m (Either [Char] [[Char]])
nameOrTupleTrg Target'
x = case Target' -> Target
forall a. WithLoc' a -> a
value' Target'
x of
        NameTrg WithLoc' VarName
x -> Either [Char] [[Char]] -> m (Either [Char] [[Char]])
forall (m :: * -> *) a. Monad m => a -> m a
return (Either [Char] [[Char]] -> m (Either [Char] [[Char]]))
-> ([Char] -> Either [Char] [[Char]])
-> [Char]
-> m (Either [Char] [[Char]])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Either [Char] [[Char]]
forall a b. a -> Either a b
Left ([Char] -> m (Either [Char] [[Char]]))
-> [Char] -> m (Either [Char] [[Char]])
forall a b. (a -> b) -> a -> b
$ VarName -> [Char]
unVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
x)
        TupleTrg [Target']
xs -> [[Char]] -> Either [Char] [[Char]]
forall a b. b -> Either a b
Right ([[Char]] -> Either [Char] [[Char]])
-> m [[Char]] -> m (Either [Char] [[Char]])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Target' -> m [Char]) -> [Target'] -> m [[Char]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Target' -> m [Char]
forall (m :: * -> *). MonadError Error m => Target' -> m [Char]
nameTrg [Target']
xs
        Target
_ -> Maybe Loc -> [Char] -> m (Either [Char] [[Char]])
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> [Char] -> m a
throwSemanticErrorAt' (Target' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Target'
x) ([Char] -> m (Either [Char] [[Char]]))
-> [Char] -> m (Either [Char] [[Char]])
forall a b. (a -> b) -> a -> b
$ [Char]
"name target or tuple target is expected, but got: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Target' -> [Char]
formatTarget Target'
x
  let nameExpr :: WithLoc' Expr -> m [Char]
nameExpr WithLoc' Expr
e = case WithLoc' Expr -> Expr
forall a. WithLoc' a -> a
value' WithLoc' Expr
e of
        Name WithLoc' VarName
x -> [Char] -> m [Char]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Char] -> m [Char]) -> [Char] -> m [Char]
forall a b. (a -> b) -> a -> b
$ VarName -> [Char]
unVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
x)
        Expr
_ -> Maybe Loc -> [Char] -> m [Char]
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> [Char] -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
e) ([Char] -> m [Char]) -> [Char] -> m [Char]
forall a b. (a -> b) -> a -> b
$ [Char]
"variable is expected, but got: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ WithLoc' Expr -> [Char]
formatExpr WithLoc' Expr
e
  case WithLoc' Expr
e of
    -- int(input())
    WithLoc' Expr
IntInput -> do
      ([Char]
x, [[Char]]
indices) <- Target' -> m ([Char], [[Char]])
forall (m :: * -> *).
MonadError Error m =>
Target' -> m ([Char], [[Char]])
subscriptTrg Target'
x
      (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]), [Statement])
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      [Statement])
forall (m :: * -> *) a. Monad m => a -> m a
return ([FormatTree] -> FormatTree
Seq [[Char] -> [[Char]] -> FormatTree
packSubscriptedVar' [Char]
x [[Char]]
indices, FormatTree
Newline], Maybe ([[Char]], Either [Char] [[Char]])
forall a. Maybe a
Nothing, [Statement]
cont)
    -- map(int, input().split())
    WithLoc' Expr
MapIntInputSplit -> do
      [([Char], [[Char]])]
outputs <- Target' -> m [([Char], [[Char]])]
forall (m :: * -> *).
MonadError Error m =>
Target' -> m [([Char], [[Char]])]
subscriptTupleTrg Target'
x
      (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]), [Statement])
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      [Statement])
forall (m :: * -> *) a. Monad m => a -> m a
return ([FormatTree] -> FormatTree
Seq ((([Char], [[Char]]) -> FormatTree)
-> [([Char], [[Char]])] -> [FormatTree]
forall a b. (a -> b) -> [a] -> [b]
map (([Char] -> [[Char]] -> FormatTree)
-> ([Char], [[Char]]) -> FormatTree
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry [Char] -> [[Char]] -> FormatTree
packSubscriptedVar') [([Char], [[Char]])]
outputs [FormatTree] -> [FormatTree] -> [FormatTree]
forall a. [a] -> [a] -> [a]
++ [FormatTree
Newline]), Maybe ([[Char]], Either [Char] [[Char]])
forall a. Maybe a
Nothing, [Statement]
cont)
    -- list(map(int, input().split()))
    WithLoc' Expr
ListMapIntInputSplit -> do
      ([Char]
x, [[Char]]
indices) <- Target' -> m ([Char], [[Char]])
forall (m :: * -> *).
MonadError Error m =>
Target' -> m ([Char], [[Char]])
subscriptTrg Target'
x
      case [Statement]
cont of
        Assert (WithLoc' Maybe Loc
_ (Compare (CallBuiltin (BuiltinLen Type
_) [WithLoc' Maybe Loc
_ (Name WithLoc' VarName
x')]) (CmpOp' CmpOp
Eq' Type
_) WithLoc' Expr
n)) : [Statement]
cont | VarName -> [Char]
unVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
x') [Char] -> [Char] -> Bool
forall a. Eq a => a -> a -> Bool
== [Char]
x -> do
          [Char]
i <- VarName -> [Char]
unVarName (VarName -> [Char])
-> (WithLoc' VarName -> VarName) -> WithLoc' VarName -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' (WithLoc' VarName -> [Char]) -> m (WithLoc' VarName) -> m [Char]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (WithLoc' VarName)
forall (m :: * -> *). MonadAlpha m => m (WithLoc' VarName)
genVarName'
          [Char]
n <- WithLoc' Expr -> m [Char]
forall (m :: * -> *).
MonadError Error m =>
WithLoc' Expr -> m [Char]
nameExpr WithLoc' Expr
n
          (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]), [Statement])
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      [Statement])
forall (m :: * -> *) a. Monad m => a -> m a
return ([FormatTree] -> FormatTree
Seq [[Char] -> FormatExpr -> FormatTree -> FormatTree
Loop [Char]
i ([Char] -> FormatExpr
Var [Char]
n) (FormatExpr -> FormatTree
Exp (FormatExpr -> [Char] -> FormatExpr
At ([Char] -> [[Char]] -> FormatExpr
packSubscriptedVar [Char]
x [[Char]]
indices) [Char]
i)), FormatTree
Newline], Maybe ([[Char]], Either [Char] [[Char]])
forall a. Maybe a
Nothing, [Statement]
cont)
        [Statement]
_ -> Maybe Loc
-> [Char]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      [Statement])
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> [Char] -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
e) [Char]
"after `xs = list(map(int, input().split()))', we need to write `assert len(xs) == n`"
    -- list(range(n))
    ListRange VarName
n -> do
      let isListRange :: Statement -> Bool
isListRange = \case
            AnnAssign Target'
_ Type
_ (ListRange VarName
n') | VarName
n' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
n -> Bool
True
            Statement
_ -> Bool
False
      [Statement]
cont <- [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Statement] -> m [Statement]) -> [Statement] -> m [Statement]
forall a b. (a -> b) -> a -> b
$ (Statement -> Bool) -> [Statement] -> [Statement]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Statement -> Bool
isListRange [Statement]
cont
      case [Statement]
cont of
        For Target'
_ (CallBuiltin Builtin
BuiltinRange1 [WithLoc' Maybe Loc
_ (Name WithLoc' VarName
n')]) [Statement]
_ : [Statement]
_ | WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
n' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
n -> (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]), [Statement])
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      [Statement])
forall (m :: * -> *) a. Monad m => a -> m a
return ([FormatTree] -> FormatTree
Seq [], Maybe ([[Char]], Either [Char] [[Char]])
forall a. Maybe a
Nothing, [Statement]
cont) -- TODO: add more strict checks
        [Statement]
_ -> Maybe Loc
-> [Char]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      [Statement])
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> [Char] -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
e) [Char]
"after some repetition of `xs = list(range(n))', we need to write `for i in range(n):`"
    -- solve(...)
    WithLoc' Maybe Loc
_ (Call (WithLoc' Maybe Loc
_ (Name (WithLoc' Maybe Loc
_ (VarName [Char]
"solve")))) [WithLoc' Expr]
args) -> do
      [[Char]]
inputs <- (WithLoc' Expr -> m [Char]) -> [WithLoc' Expr] -> m [[Char]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM WithLoc' Expr -> m [Char]
forall (m :: * -> *).
MonadError Error m =>
WithLoc' Expr -> m [Char]
nameExpr [WithLoc' Expr]
args
      Either [Char] [[Char]]
output <- Target' -> m (Either [Char] [[Char]])
forall (m :: * -> *).
MonadError Error m =>
Target' -> m (Either [Char] [[Char]])
nameOrTupleTrg Target'
x
      (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]), [Statement])
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      [Statement])
forall (m :: * -> *) a. Monad m => a -> m a
return ([FormatTree] -> FormatTree
Seq [], ([[Char]], Either [Char] [[Char]])
-> Maybe ([[Char]], Either [Char] [[Char]])
forall a. a -> Maybe a
Just ([[Char]]
inputs, Either [Char] [[Char]]
output), [Statement]
cont)
    WithLoc' Expr
_ -> Maybe Loc
-> [Char]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      [Statement])
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> [Char] -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
e) [Char]
"assignments in main function must be `x = int(input())', `x, y, z = map(int, input().split())', `xs = list(map(int, input().split()))', `xs = list(range(n))' or `x, y, z = solve(a, b, c)'"

parseFor :: MonadError Error m => ([Statement] -> m (FormatTree, Maybe ([String], Either String [String]), FormatTree)) -> Target' -> Expr' -> [Statement] -> m (FormatTree, FormatTree)
parseFor :: ([Statement]
 -> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
       FormatTree))
-> Target'
-> WithLoc' Expr
-> [Statement]
-> m (FormatTree, FormatTree)
parseFor [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
go Target'
x WithLoc' Expr
e [Statement]
body = do
  WithLoc' VarName
x <- case Target' -> Target
forall a. WithLoc' a -> a
value' Target'
x of
    NameTrg WithLoc' VarName
x -> WithLoc' VarName -> m (WithLoc' VarName)
forall (m :: * -> *) a. Monad m => a -> m a
return WithLoc' VarName
x
    Target
_ -> Maybe Loc -> [Char] -> m (WithLoc' VarName)
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> [Char] -> m a
throwSemanticErrorAt' (Target' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Target'
x) ([Char] -> m (WithLoc' VarName)) -> [Char] -> m (WithLoc' VarName)
forall a b. (a -> b) -> a -> b
$ [Char]
"for loops in main function must use `range' like `for i in range(n): ...'" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Target' -> [Char]
formatTarget Target'
x
  WithLoc' Expr
n <- case WithLoc' Expr
e of
    CallBuiltin Builtin
BuiltinRange1 [WithLoc' Expr
n] -> WithLoc' Expr -> m (WithLoc' Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return WithLoc' Expr
n
    WithLoc' Expr
_ -> Maybe Loc -> [Char] -> m (WithLoc' Expr)
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> [Char] -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
e) ([Char] -> m (WithLoc' Expr)) -> [Char] -> m (WithLoc' Expr)
forall a b. (a -> b) -> a -> b
$ [Char]
"for loops in main function must use `range' like `for i in range(n): ...': " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ WithLoc' Expr -> [Char]
formatExpr WithLoc' Expr
e
  Either (WithLoc' VarName) (WithLoc' VarName, Integer)
n <- case WithLoc' Expr -> Expr
forall a. WithLoc' a -> a
value' WithLoc' Expr
n of
    Name WithLoc' VarName
n -> Either (WithLoc' VarName) (WithLoc' VarName, Integer)
-> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer))
forall (m :: * -> *) a. Monad m => a -> m a
return (Either (WithLoc' VarName) (WithLoc' VarName, Integer)
 -> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer)))
-> Either (WithLoc' VarName) (WithLoc' VarName, Integer)
-> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer))
forall a b. (a -> b) -> a -> b
$ (WithLoc' VarName, Integer)
-> Either (WithLoc' VarName) (WithLoc' VarName, Integer)
forall a b. b -> Either a b
Right (WithLoc' VarName
n, Integer
0)
    BinOp (WithLoc' Maybe Loc
_ (Name WithLoc' VarName
n)) Operator
Add (WithLoc' Maybe Loc
_ (Constant (ConstInt Integer
k))) -> Either (WithLoc' VarName) (WithLoc' VarName, Integer)
-> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer))
forall (m :: * -> *) a. Monad m => a -> m a
return (Either (WithLoc' VarName) (WithLoc' VarName, Integer)
 -> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer)))
-> Either (WithLoc' VarName) (WithLoc' VarName, Integer)
-> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer))
forall a b. (a -> b) -> a -> b
$ (WithLoc' VarName, Integer)
-> Either (WithLoc' VarName) (WithLoc' VarName, Integer)
forall a b. b -> Either a b
Right (WithLoc' VarName
n, Integer
k)
    BinOp (WithLoc' Maybe Loc
_ (Name WithLoc' VarName
n)) Operator
Sub (WithLoc' Maybe Loc
_ (Constant (ConstInt Integer
k))) -> Either (WithLoc' VarName) (WithLoc' VarName, Integer)
-> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer))
forall (m :: * -> *) a. Monad m => a -> m a
return (Either (WithLoc' VarName) (WithLoc' VarName, Integer)
 -> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer)))
-> Either (WithLoc' VarName) (WithLoc' VarName, Integer)
-> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer))
forall a b. (a -> b) -> a -> b
$ (WithLoc' VarName, Integer)
-> Either (WithLoc' VarName) (WithLoc' VarName, Integer)
forall a b. b -> Either a b
Right (WithLoc' VarName
n, - Integer
k)
    Call (WithLoc' Maybe Loc
_ (Constant (ConstBuiltin (BuiltinLen Type
_)))) [WithLoc' Maybe Loc
_ (Name WithLoc' VarName
xs)] -> Either (WithLoc' VarName) (WithLoc' VarName, Integer)
-> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer))
forall (m :: * -> *) a. Monad m => a -> m a
return (Either (WithLoc' VarName) (WithLoc' VarName, Integer)
 -> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer)))
-> Either (WithLoc' VarName) (WithLoc' VarName, Integer)
-> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer))
forall a b. (a -> b) -> a -> b
$ WithLoc' VarName
-> Either (WithLoc' VarName) (WithLoc' VarName, Integer)
forall a b. a -> Either a b
Left WithLoc' VarName
xs
    Expr
_ -> Maybe Loc
-> [Char]
-> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer))
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> [Char] -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
n) ([Char]
 -> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer)))
-> [Char]
-> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer))
forall a b. (a -> b) -> a -> b
$ [Char]
"for loops in main function must use `range(x)', `range(x + k)', `range(x - k)', `range(len(xs))`: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ WithLoc' Expr -> [Char]
formatExpr WithLoc' Expr
n
  FormatExpr
n <- FormatExpr -> m FormatExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (FormatExpr -> m FormatExpr) -> FormatExpr -> m FormatExpr
forall a b. (a -> b) -> a -> b
$ case Either (WithLoc' VarName) (WithLoc' VarName, Integer)
n of
    Right (WithLoc' VarName
n, Integer
k) ->
      let n' :: FormatExpr
n' = [Char] -> FormatExpr
Var (VarName -> [Char]
unVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
n))
       in if Integer
k Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 then FormatExpr
n' else FormatExpr -> Integer -> FormatExpr
Plus FormatExpr
n' Integer
k
    Left WithLoc' VarName
xs -> FormatExpr -> FormatExpr
Len ([Char] -> FormatExpr
Var (VarName -> [Char]
unVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
xs)))
  (FormatTree
input, Maybe ([[Char]], Either [Char] [[Char]])
solve, FormatTree
output) <- [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
go [Statement]
body
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe ([[Char]], Either [Char] [[Char]]) -> Bool
forall a. Maybe a -> Bool
isJust Maybe ([[Char]], Either [Char] [[Char]])
solve) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    [Char] -> m ()
forall (m :: * -> *) a. MonadError Error m => [Char] -> m a
throwSemanticError [Char]
"cannot call `solve(...)' in for loop"
  let x' :: [Char]
x' = VarName -> [Char]
unVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
x)
  (FormatTree, FormatTree) -> m (FormatTree, FormatTree)
forall (m :: * -> *) a. Monad m => a -> m a
return ([Char] -> FormatExpr -> FormatTree -> FormatTree
Loop [Char]
x' FormatExpr
n FormatTree
input, [Char] -> FormatExpr -> FormatTree -> FormatTree
Loop [Char]
x' FormatExpr
n FormatTree
output)

parseExprStatement :: (MonadAlpha m, MonadError Error m) => Expr' -> m FormatTree
parseExprStatement :: WithLoc' Expr -> m FormatTree
parseExprStatement WithLoc' Expr
e = do
  let subscriptExpr :: WithLoc' Expr -> m ([Char], [[Char]])
subscriptExpr WithLoc' Expr
e = case WithLoc' Expr -> Expr
forall a. WithLoc' a -> a
value' WithLoc' Expr
e of
        Name WithLoc' VarName
x -> ([Char], [[Char]]) -> m ([Char], [[Char]])
forall (m :: * -> *) a. Monad m => a -> m a
return (VarName -> [Char]
unVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
x), [])
        Subscript WithLoc' Expr
e (WithLoc' Maybe Loc
_ (Name WithLoc' VarName
i)) -> ([[Char]] -> [[Char]]) -> ([Char], [[Char]]) -> ([Char], [[Char]])
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ([[Char]] -> [[Char]] -> [[Char]]
forall a. [a] -> [a] -> [a]
++ [VarName -> [Char]
unVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
i)]) (([Char], [[Char]]) -> ([Char], [[Char]]))
-> m ([Char], [[Char]]) -> m ([Char], [[Char]])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WithLoc' Expr -> m ([Char], [[Char]])
subscriptExpr WithLoc' Expr
e
        Expr
_ -> Maybe Loc -> [Char] -> m ([Char], [[Char]])
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> [Char] -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
e) ([Char] -> m ([Char], [[Char]])) -> [Char] -> m ([Char], [[Char]])
forall a b. (a -> b) -> a -> b
$ [Char]
"subscripted variable is expected, but got: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ WithLoc' Expr -> [Char]
formatExpr WithLoc' Expr
e
  let starredExpr :: WithLoc' Expr -> m ([Char], [[Char]], Bool)
starredExpr WithLoc' Expr
e = do
        (WithLoc' Expr
e, Bool
starred) <- (WithLoc' Expr, Bool) -> m (WithLoc' Expr, Bool)
forall (m :: * -> *) a. Monad m => a -> m a
return ((WithLoc' Expr, Bool) -> m (WithLoc' Expr, Bool))
-> (WithLoc' Expr, Bool) -> m (WithLoc' Expr, Bool)
forall a b. (a -> b) -> a -> b
$ case WithLoc' Expr -> Expr
forall a. WithLoc' a -> a
value' WithLoc' Expr
e of
          Starred WithLoc' Expr
e -> (WithLoc' Expr
e, Bool
True)
          Expr
_ -> (WithLoc' Expr
e, Bool
False)
        ([Char]
x, [[Char]]
indices) <- WithLoc' Expr -> m ([Char], [[Char]])
forall (m :: * -> *).
MonadError Error m =>
WithLoc' Expr -> m ([Char], [[Char]])
subscriptExpr WithLoc' Expr
e
        ([Char], [[Char]], Bool) -> m ([Char], [[Char]], Bool)
forall (m :: * -> *) a. Monad m => a -> m a
return ([Char]
x, [[Char]]
indices, Bool
starred)
  let pack :: ([Char], [[Char]], Bool) -> m FormatTree
pack ([Char]
x, [[Char]]
indices, Bool
starred)
        | Bool -> Bool
not Bool
starred = FormatTree -> m FormatTree
forall (m :: * -> *) a. Monad m => a -> m a
return (FormatTree -> m FormatTree) -> FormatTree -> m FormatTree
forall a b. (a -> b) -> a -> b
$ [Char] -> [[Char]] -> FormatTree
packSubscriptedVar' [Char]
x [[Char]]
indices
        | Bool
otherwise = do
          let xs :: FormatExpr
xs = [Char] -> [[Char]] -> FormatExpr
packSubscriptedVar [Char]
x [[Char]]
indices
          [Char]
i <- VarName -> [Char]
unVarName (VarName -> [Char])
-> (WithLoc' VarName -> VarName) -> WithLoc' VarName -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' (WithLoc' VarName -> [Char]) -> m (WithLoc' VarName) -> m [Char]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (WithLoc' VarName)
forall (m :: * -> *). MonadAlpha m => m (WithLoc' VarName)
genVarName'
          FormatTree -> m FormatTree
forall (m :: * -> *) a. Monad m => a -> m a
return (FormatTree -> m FormatTree) -> FormatTree -> m FormatTree
forall a b. (a -> b) -> a -> b
$ [Char] -> FormatExpr -> FormatTree -> FormatTree
Loop [Char]
i (FormatExpr -> FormatExpr
Len FormatExpr
xs) ([Char] -> [[Char]] -> FormatTree
packSubscriptedVar' [Char]
x ([[Char]]
indices [[Char]] -> [[Char]] -> [[Char]]
forall a. [a] -> [a] -> [a]
++ [[Char]
i]))
  case WithLoc' Expr
e of
    CallBuiltin (BuiltinPrint [Type]
_) [WithLoc' Expr]
args -> do
      [([Char], [[Char]], Bool)]
outputs <- (WithLoc' Expr -> m ([Char], [[Char]], Bool))
-> [WithLoc' Expr] -> m [([Char], [[Char]], Bool)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM WithLoc' Expr -> m ([Char], [[Char]], Bool)
forall (m :: * -> *).
MonadError Error m =>
WithLoc' Expr -> m ([Char], [[Char]], Bool)
starredExpr [WithLoc' Expr]
args
      [FormatTree]
outputs <- (([Char], [[Char]], Bool) -> m FormatTree)
-> [([Char], [[Char]], Bool)] -> m [FormatTree]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char], [[Char]], Bool) -> m FormatTree
forall (m :: * -> *).
MonadAlpha m =>
([Char], [[Char]], Bool) -> m FormatTree
pack [([Char], [[Char]], Bool)]
outputs
      FormatTree -> m FormatTree
forall (m :: * -> *) a. Monad m => a -> m a
return (FormatTree -> m FormatTree) -> FormatTree -> m FormatTree
forall a b. (a -> b) -> a -> b
$ [FormatTree] -> FormatTree
Seq ([FormatTree]
outputs [FormatTree] -> [FormatTree] -> [FormatTree]
forall a. [a] -> [a] -> [a]
++ [FormatTree
Newline])
    WithLoc' Expr
_ -> Maybe Loc -> [Char] -> m FormatTree
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> [Char] -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
e) [Char]
"only `print(...)' is allowed for expr statements in main function"

parseMain :: (MonadAlpha m, MonadError Error m) => MainFunction -> m IOFormat
parseMain :: MainFunction -> m IOFormat
parseMain (Maybe Loc
loc, [(WithLoc' VarName, Type)]
_, Type
_, [Statement]
body) = Maybe Loc -> m IOFormat -> m IOFormat
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> m a -> m a
wrapAt' Maybe Loc
loc (m IOFormat -> m IOFormat) -> m IOFormat -> m IOFormat
forall a b. (a -> b) -> a -> b
$ (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]), FormatTree)
-> m IOFormat
forall (m :: * -> *).
MonadError Error m =>
(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]), FormatTree)
-> m IOFormat
pack ((FormatTree, Maybe ([[Char]], Either [Char] [[Char]]), FormatTree)
 -> m IOFormat)
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
-> m IOFormat
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
go [] [Statement]
body
  where
    pack :: MonadError Error m => (FormatTree, Maybe ([String], Either String [String]), FormatTree) -> m IOFormat
    pack :: (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]), FormatTree)
-> m IOFormat
pack (FormatTree
_, Maybe ([[Char]], Either [Char] [[Char]])
Nothing, FormatTree
_) = [Char] -> m IOFormat
forall (m :: * -> *) a. MonadError Error m => [Char] -> m a
throwSemanticError [Char]
"main function must call solve function"
    pack (FormatTree
inputTree, Just ([[Char]]
inputVariables, Either [Char] [[Char]]
outputVariables), FormatTree
outputTree) =
      IOFormat -> m IOFormat
forall (m :: * -> *) a. Monad m => a -> m a
return (IOFormat -> m IOFormat) -> IOFormat -> m IOFormat
forall a b. (a -> b) -> a -> b
$
        IOFormat :: [[Char]]
-> FormatTree -> Either [Char] [[Char]] -> FormatTree -> IOFormat
IOFormat
          { inputTree :: FormatTree
inputTree = FormatTree
inputTree,
            inputVariables :: [[Char]]
inputVariables = [[Char]]
inputVariables,
            outputVariables :: Either [Char] [[Char]]
outputVariables = Either [Char] [[Char]]
outputVariables,
            outputTree :: FormatTree
outputTree = FormatTree
outputTree
          }
    go :: (MonadAlpha m, MonadError Error m) => [(FormatTree, Maybe ([String], Either String [String]), FormatTree)] -> [Statement] -> m (FormatTree, Maybe ([String], Either String [String]), FormatTree)
    go :: [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
go [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
formats = \case
      Return WithLoc' Expr
_ : [Statement]
_ -> [Char]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
forall (m :: * -> *) a. MonadError Error m => [Char] -> m a
throwSemanticError [Char]
"return statement is not allowd in main function"
      AugAssign Target'
_ Operator
_ WithLoc' Expr
_ : [Statement]
_ -> [Char]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
forall (m :: * -> *) a. MonadError Error m => [Char] -> m a
throwSemanticError [Char]
"augumented assignment statement is not allowd in main function"
      AnnAssign Target'
x Type
t WithLoc' Expr
e : [Statement]
cont -> do
        (FormatTree
inputs, Maybe ([[Char]], Either [Char] [[Char]])
solve, [Statement]
cont) <- Target'
-> Type
-> WithLoc' Expr
-> [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      [Statement])
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
Target'
-> Type
-> WithLoc' Expr
-> [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      [Statement])
parseAnnAssign Target'
x Type
t WithLoc' Expr
e [Statement]
cont
        [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
go ([(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
formats [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
-> [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
     FormatTree)]
-> [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
     FormatTree)]
forall a. [a] -> [a] -> [a]
++ [(FormatTree
inputs, Maybe ([[Char]], Either [Char] [[Char]])
solve, [FormatTree] -> FormatTree
Seq [])]) [Statement]
cont
      For Target'
x WithLoc' Expr
e [Statement]
body : [Statement]
cont -> do
        (FormatTree
inputs, FormatTree
outputs) <- ([Statement]
 -> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
       FormatTree))
-> Target'
-> WithLoc' Expr
-> [Statement]
-> m (FormatTree, FormatTree)
forall (m :: * -> *).
MonadError Error m =>
([Statement]
 -> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
       FormatTree))
-> Target'
-> WithLoc' Expr
-> [Statement]
-> m (FormatTree, FormatTree)
parseFor ([(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
go []) Target'
x WithLoc' Expr
e [Statement]
body
        [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
go ([(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
formats [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
-> [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
     FormatTree)]
-> [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
     FormatTree)]
forall a. [a] -> [a] -> [a]
++ [(FormatTree
inputs, Maybe ([[Char]], Either [Char] [[Char]])
forall a. Maybe a
Nothing, FormatTree
outputs)]) [Statement]
cont
      If WithLoc' Expr
_ [Statement]
_ [Statement]
_ : [Statement]
_ -> [Char]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
forall (m :: * -> *) a. MonadError Error m => [Char] -> m a
throwSemanticError [Char]
"if statement is not allowd in main function"
      Assert WithLoc' Expr
_ : [Statement]
_ -> [Char]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
forall (m :: * -> *) a. MonadError Error m => [Char] -> m a
throwSemanticError [Char]
"assert statement is allowd only after `xs = list(map(int, input().split()))` in main function"
      Expr' WithLoc' Expr
e : [Statement]
cont -> do
        FormatTree
output <- WithLoc' Expr -> m FormatTree
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
WithLoc' Expr -> m FormatTree
parseExprStatement WithLoc' Expr
e
        [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
go ([(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
formats [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
-> [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
     FormatTree)]
-> [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
     FormatTree)]
forall a. [a] -> [a] -> [a]
++ [([FormatTree] -> FormatTree
Seq [], Maybe ([[Char]], Either [Char] [[Char]])
forall a. Maybe a
Nothing, FormatTree
output)]) [Statement]
cont
      [] -> do
        let input :: FormatTree
input = [FormatTree] -> FormatTree
Seq (((FormatTree, Maybe ([[Char]], Either [Char] [[Char]]), FormatTree)
 -> FormatTree)
-> [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
     FormatTree)]
-> [FormatTree]
forall a b. (a -> b) -> [a] -> [b]
map (\(FormatTree
x, Maybe ([[Char]], Either [Char] [[Char]])
_, FormatTree
_) -> FormatTree
x) [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
formats)
        let outputs :: FormatTree
outputs = [FormatTree] -> FormatTree
Seq (((FormatTree, Maybe ([[Char]], Either [Char] [[Char]]), FormatTree)
 -> FormatTree)
-> [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
     FormatTree)]
-> [FormatTree]
forall a b. (a -> b) -> [a] -> [b]
map (\(FormatTree
_, Maybe ([[Char]], Either [Char] [[Char]])
_, FormatTree
z) -> FormatTree
z) [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
formats)
        Maybe ([[Char]], Either [Char] [[Char]])
solve <- case ((FormatTree, Maybe ([[Char]], Either [Char] [[Char]]), FormatTree)
 -> Maybe ([[Char]], Either [Char] [[Char]]))
-> [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
     FormatTree)]
-> [([[Char]], Either [Char] [[Char]])]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(FormatTree
_, Maybe ([[Char]], Either [Char] [[Char]])
y, FormatTree
_) -> Maybe ([[Char]], Either [Char] [[Char]])
y) [(FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
  FormatTree)]
formats of
          [] -> Maybe ([[Char]], Either [Char] [[Char]])
-> m (Maybe ([[Char]], Either [Char] [[Char]]))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ([[Char]], Either [Char] [[Char]])
forall a. Maybe a
Nothing
          [([[Char]], Either [Char] [[Char]])
solve] -> Maybe ([[Char]], Either [Char] [[Char]])
-> m (Maybe ([[Char]], Either [Char] [[Char]]))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ([[Char]], Either [Char] [[Char]])
 -> m (Maybe ([[Char]], Either [Char] [[Char]])))
-> Maybe ([[Char]], Either [Char] [[Char]])
-> m (Maybe ([[Char]], Either [Char] [[Char]]))
forall a b. (a -> b) -> a -> b
$ ([[Char]], Either [Char] [[Char]])
-> Maybe ([[Char]], Either [Char] [[Char]])
forall a. a -> Maybe a
Just ([[Char]], Either [Char] [[Char]])
solve
          [([[Char]], Either [Char] [[Char]])]
_ -> [Char] -> m (Maybe ([[Char]], Either [Char] [[Char]]))
forall (m :: * -> *) a. MonadError Error m => [Char] -> m a
throwSemanticError [Char]
"cannot call solve function twice"
        (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]), FormatTree)
-> m (FormatTree, Maybe ([[Char]], Either [Char] [[Char]]),
      FormatTree)
forall (m :: * -> *) a. Monad m => a -> m a
return (FormatTree
input, Maybe ([[Char]], Either [Char] [[Char]])
solve, FormatTree
outputs)

run :: (MonadAlpha m, MonadError Error m) => Program -> m (Maybe IOFormat, Program)
run :: Program -> m (Maybe IOFormat, Program)
run Program
prog = [Char]
-> m (Maybe IOFormat, Program) -> m (Maybe IOFormat, Program)
forall (m :: * -> *) a. MonadError Error m => [Char] -> m a -> m a
wrapError' [Char]
"Jikka.RestrictedPython.Convert.ParseMain" (m (Maybe IOFormat, Program) -> m (Maybe IOFormat, Program))
-> m (Maybe IOFormat, Program) -> m (Maybe IOFormat, Program)
forall a b. (a -> b) -> a -> b
$ do
  (Maybe MainFunction
main, Program
prog) <- (Maybe MainFunction, Program) -> m (Maybe MainFunction, Program)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Maybe MainFunction, Program) -> m (Maybe MainFunction, Program))
-> (Maybe MainFunction, Program) -> m (Maybe MainFunction, Program)
forall a b. (a -> b) -> a -> b
$ Program -> (Maybe MainFunction, Program)
splitMain Program
prog
  Maybe IOFormat
main <- Maybe MainFunction
-> (MainFunction -> m IOFormat) -> m (Maybe IOFormat)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Maybe MainFunction
main ((MainFunction -> m IOFormat) -> m (Maybe IOFormat))
-> (MainFunction -> m IOFormat) -> m (Maybe IOFormat)
forall a b. (a -> b) -> a -> b
$ \MainFunction
main -> do
    MainFunction -> m ()
forall (m :: * -> *). MonadError Error m => MainFunction -> m ()
checkMainType MainFunction
main
    IOFormat
main <- MainFunction -> m IOFormat
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
MainFunction -> m IOFormat
parseMain MainFunction
main
    IOFormat -> m IOFormat
forall (m :: * -> *) a. Monad m => a -> m a
return (IOFormat -> m IOFormat) -> IOFormat -> m IOFormat
forall a b. (a -> b) -> a -> b
$ IOFormat -> IOFormat
normalizeIOFormat IOFormat
main
  (Maybe IOFormat, Program) -> m (Maybe IOFormat, Program)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe IOFormat
main, Program
prog)