{-# 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 (Just String
"main") Maybe Int
Nothing Maybe NameHint
Nothing)) [(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)]
_ -> String -> m ()
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwTypeError String
"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
_ -> String -> m ()
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwTypeError String
"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 ([String], Either String [String]),
      [Statement])
parseAnnAssign Target'
x Type
_ WithLoc' Expr
e [Statement]
cont = do
  let subscriptTrg :: Target' -> m (String, [String])
subscriptTrg Target'
x = case Target' -> Target
forall a. WithLoc' a -> a
value' Target'
x of
        NameTrg WithLoc' VarName
x -> (String, [String]) -> m (String, [String])
forall (m :: * -> *) a. Monad m => a -> m a
return (VarName -> String
formatVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
x), [])
        SubscriptTrg Target'
x (WithLoc' Maybe Loc
_ (Name WithLoc' VarName
i)) -> ([String] -> [String]) -> (String, [String]) -> (String, [String])
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ([String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ [VarName -> String
formatVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
i)]) ((String, [String]) -> (String, [String]))
-> m (String, [String]) -> m (String, [String])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Target' -> m (String, [String])
subscriptTrg Target'
x
        Target
_ -> Maybe Loc -> String -> m (String, [String])
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> String -> m a
throwSemanticErrorAt' (Target' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Target'
x) (String -> m (String, [String])) -> String -> m (String, [String])
forall a b. (a -> b) -> a -> b
$ String
"name target or subscript target is expected, but got: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Target' -> String
formatTarget Target'
x
  let subscriptTupleTrg :: Target' -> m [(String, [String])]
subscriptTupleTrg Target'
x = case Target' -> Target
forall a. WithLoc' a -> a
value' Target'
x of
        TupleTrg [Target']
xs -> (Target' -> m (String, [String]))
-> [Target'] -> m [(String, [String])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Target' -> m (String, [String])
forall (m :: * -> *).
MonadError Error m =>
Target' -> m (String, [String])
subscriptTrg [Target']
xs
        Target
_ -> Maybe Loc -> String -> m [(String, [String])]
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> String -> m a
throwSemanticErrorAt' (Target' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Target'
x) (String -> m [(String, [String])])
-> String -> m [(String, [String])]
forall a b. (a -> b) -> a -> b
$ String
"tuple target is expected, but got: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Target' -> String
formatTarget Target'
x
  let nameTrg :: Target' -> m String
nameTrg Target'
x = case Target' -> Target
forall a. WithLoc' a -> a
value' Target'
x of
        NameTrg WithLoc' VarName
x -> String -> m String
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> m String) -> String -> m String
forall a b. (a -> b) -> a -> b
$ VarName -> String
formatVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
x)
        Target
_ -> Maybe Loc -> String -> m String
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> String -> m a
throwSemanticErrorAt' (Target' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Target'
x) (String -> m String) -> String -> m String
forall a b. (a -> b) -> a -> b
$ String
"name target is expected, but got: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Target' -> String
formatTarget Target'
x
  let nameOrTupleTrg :: Target' -> m (Either String [String])
nameOrTupleTrg Target'
x = case Target' -> Target
forall a. WithLoc' a -> a
value' Target'
x of
        NameTrg WithLoc' VarName
x -> Either String [String] -> m (Either String [String])
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String [String] -> m (Either String [String]))
-> (String -> Either String [String])
-> String
-> m (Either String [String])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Either String [String]
forall a b. a -> Either a b
Left (String -> m (Either String [String]))
-> String -> m (Either String [String])
forall a b. (a -> b) -> a -> b
$ VarName -> String
formatVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
x)
        TupleTrg [Target']
xs -> [String] -> Either String [String]
forall a b. b -> Either a b
Right ([String] -> Either String [String])
-> m [String] -> m (Either String [String])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Target' -> m String) -> [Target'] -> m [String]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Target' -> m String
forall (m :: * -> *). MonadError Error m => Target' -> m String
nameTrg [Target']
xs
        Target
_ -> Maybe Loc -> String -> m (Either String [String])
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> String -> m a
throwSemanticErrorAt' (Target' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Target'
x) (String -> m (Either String [String]))
-> String -> m (Either String [String])
forall a b. (a -> b) -> a -> b
$ String
"name target or tuple target is expected, but got: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Target' -> String
formatTarget Target'
x
  let nameExpr :: WithLoc' Expr -> m String
nameExpr WithLoc' Expr
e = case WithLoc' Expr -> Expr
forall a. WithLoc' a -> a
value' WithLoc' Expr
e of
        Name WithLoc' VarName
x -> String -> m String
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> m String) -> String -> m String
forall a b. (a -> b) -> a -> b
$ VarName -> String
formatVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
x)
        Expr
_ -> Maybe Loc -> String -> m String
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> String -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
e) (String -> m String) -> String -> m String
forall a b. (a -> b) -> a -> b
$ String
"variable is expected, but got: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ WithLoc' Expr -> String
formatExpr WithLoc' Expr
e
  case WithLoc' Expr
e of
    -- int(input())
    WithLoc' Expr
IntInput -> do
      (String
x, [String]
indices) <- Target' -> m (String, [String])
forall (m :: * -> *).
MonadError Error m =>
Target' -> m (String, [String])
subscriptTrg Target'
x
      (FormatTree, Maybe ([String], Either String [String]), [Statement])
-> m (FormatTree, Maybe ([String], Either String [String]),
      [Statement])
forall (m :: * -> *) a. Monad m => a -> m a
return ([FormatTree] -> FormatTree
Seq [String -> [String] -> FormatTree
packSubscriptedVar' String
x [String]
indices, FormatTree
Newline], Maybe ([String], Either String [String])
forall a. Maybe a
Nothing, [Statement]
cont)
    -- map(int, input().split())
    WithLoc' Expr
MapIntInputSplit -> do
      [(String, [String])]
outputs <- Target' -> m [(String, [String])]
forall (m :: * -> *).
MonadError Error m =>
Target' -> m [(String, [String])]
subscriptTupleTrg Target'
x
      (FormatTree, Maybe ([String], Either String [String]), [Statement])
-> m (FormatTree, Maybe ([String], Either String [String]),
      [Statement])
forall (m :: * -> *) a. Monad m => a -> m a
return ([FormatTree] -> FormatTree
Seq (((String, [String]) -> FormatTree)
-> [(String, [String])] -> [FormatTree]
forall a b. (a -> b) -> [a] -> [b]
map ((String -> [String] -> FormatTree)
-> (String, [String]) -> FormatTree
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry String -> [String] -> FormatTree
packSubscriptedVar') [(String, [String])]
outputs [FormatTree] -> [FormatTree] -> [FormatTree]
forall a. [a] -> [a] -> [a]
++ [FormatTree
Newline]), Maybe ([String], Either String [String])
forall a. Maybe a
Nothing, [Statement]
cont)
    -- list(map(int, input().split()))
    WithLoc' Expr
ListMapIntInputSplit -> do
      (String
x, [String]
indices) <- Target' -> m (String, [String])
forall (m :: * -> *).
MonadError Error m =>
Target' -> m (String, [String])
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 -> String
formatVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
x') String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
x -> do
          String
i <- VarName -> String
formatVarName (VarName -> String)
-> (WithLoc' VarName -> VarName) -> WithLoc' VarName -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' (WithLoc' VarName -> String) -> m (WithLoc' VarName) -> m String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (WithLoc' VarName)
forall (m :: * -> *). MonadAlpha m => m (WithLoc' VarName)
genVarName'
          String
n <- WithLoc' Expr -> m String
forall (m :: * -> *).
MonadError Error m =>
WithLoc' Expr -> m String
nameExpr WithLoc' Expr
n
          (FormatTree, Maybe ([String], Either String [String]), [Statement])
-> m (FormatTree, Maybe ([String], Either String [String]),
      [Statement])
forall (m :: * -> *) a. Monad m => a -> m a
return ([FormatTree] -> FormatTree
Seq [String -> FormatExpr -> FormatTree -> FormatTree
Loop String
i (String -> FormatExpr
Var String
n) (FormatExpr -> FormatTree
Exp (FormatExpr -> String -> FormatExpr
At (String -> [String] -> FormatExpr
packSubscriptedVar String
x [String]
indices) String
i)), FormatTree
Newline], Maybe ([String], Either String [String])
forall a. Maybe a
Nothing, [Statement]
cont)
        [Statement]
_ -> Maybe Loc
-> String
-> m (FormatTree, Maybe ([String], Either String [String]),
      [Statement])
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> String -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
e) String
"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 ([String], Either String [String]), [Statement])
-> m (FormatTree, Maybe ([String], Either String [String]),
      [Statement])
forall (m :: * -> *) a. Monad m => a -> m a
return ([FormatTree] -> FormatTree
Seq [], Maybe ([String], Either String [String])
forall a. Maybe a
Nothing, [Statement]
cont) -- TODO: add more strict checks
        [Statement]
_ -> Maybe Loc
-> String
-> m (FormatTree, Maybe ([String], Either String [String]),
      [Statement])
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> String -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
e) String
"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 (Just String
"solve") Maybe Int
Nothing Maybe NameHint
Nothing)))) [WithLoc' Expr]
args) -> do
      [String]
inputs <- (WithLoc' Expr -> m String) -> [WithLoc' Expr] -> m [String]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM WithLoc' Expr -> m String
forall (m :: * -> *).
MonadError Error m =>
WithLoc' Expr -> m String
nameExpr [WithLoc' Expr]
args
      Either String [String]
output <- Target' -> m (Either String [String])
forall (m :: * -> *).
MonadError Error m =>
Target' -> m (Either String [String])
nameOrTupleTrg Target'
x
      (FormatTree, Maybe ([String], Either String [String]), [Statement])
-> m (FormatTree, Maybe ([String], Either String [String]),
      [Statement])
forall (m :: * -> *) a. Monad m => a -> m a
return ([FormatTree] -> FormatTree
Seq [], ([String], Either String [String])
-> Maybe ([String], Either String [String])
forall a. a -> Maybe a
Just ([String]
inputs, Either String [String]
output), [Statement]
cont)
    WithLoc' Expr
_ -> Maybe Loc
-> String
-> m (FormatTree, Maybe ([String], Either String [String]),
      [Statement])
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> String -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
e) String
"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 ([String], Either String [String]),
       FormatTree))
-> Target'
-> WithLoc' Expr
-> [Statement]
-> m (FormatTree, FormatTree)
parseFor [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      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 -> String -> m (WithLoc' VarName)
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> String -> m a
throwSemanticErrorAt' (Target' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Target'
x) (String -> m (WithLoc' VarName)) -> String -> m (WithLoc' VarName)
forall a b. (a -> b) -> a -> b
$ String
"for loops in main function must use `range' like `for i in range(n): ...'" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Target' -> String
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 -> String -> m (WithLoc' Expr)
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> String -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
e) (String -> m (WithLoc' Expr)) -> String -> m (WithLoc' Expr)
forall a b. (a -> b) -> a -> b
$ String
"for loops in main function must use `range' like `for i in range(n): ...': " String -> String -> String
forall a. [a] -> [a] -> [a]
++ WithLoc' Expr -> String
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
-> String
-> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer))
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> String -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
n) (String
 -> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer)))
-> String
-> m (Either (WithLoc' VarName) (WithLoc' VarName, Integer))
forall a b. (a -> b) -> a -> b
$ String
"for loops in main function must use `range(x)', `range(x + k)', `range(x - k)', `range(len(xs))`: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ WithLoc' Expr -> String
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' = String -> FormatExpr
Var (VarName -> String
formatVarName (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 (String -> FormatExpr
Var (VarName -> String
formatVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
xs)))
  (FormatTree
input, Maybe ([String], Either String [String])
solve, FormatTree
output) <- [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
go [Statement]
body
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe ([String], Either String [String]) -> Bool
forall a. Maybe a -> Bool
isJust Maybe ([String], Either String [String])
solve) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    String -> m ()
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwSemanticError String
"cannot call `solve(...)' in for loop"
  let x' :: String
x' = VarName -> String
formatVarName (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 (String -> FormatExpr -> FormatTree -> FormatTree
Loop String
x' FormatExpr
n FormatTree
input, String -> FormatExpr -> FormatTree -> FormatTree
Loop String
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 (String, [String])
subscriptExpr WithLoc' Expr
e = case WithLoc' Expr -> Expr
forall a. WithLoc' a -> a
value' WithLoc' Expr
e of
        Name WithLoc' VarName
x -> (String, [String]) -> m (String, [String])
forall (m :: * -> *) a. Monad m => a -> m a
return (VarName -> String
formatVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
x), [])
        Subscript WithLoc' Expr
e (WithLoc' Maybe Loc
_ (Name WithLoc' VarName
i)) -> ([String] -> [String]) -> (String, [String]) -> (String, [String])
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ([String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ [VarName -> String
formatVarName (WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' WithLoc' VarName
i)]) ((String, [String]) -> (String, [String]))
-> m (String, [String]) -> m (String, [String])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WithLoc' Expr -> m (String, [String])
subscriptExpr WithLoc' Expr
e
        Expr
_ -> Maybe Loc -> String -> m (String, [String])
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> String -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
e) (String -> m (String, [String])) -> String -> m (String, [String])
forall a b. (a -> b) -> a -> b
$ String
"subscripted variable is expected, but got: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ WithLoc' Expr -> String
formatExpr WithLoc' Expr
e
  let starredExpr :: WithLoc' Expr -> m (String, [String], 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)
        (String
x, [String]
indices) <- WithLoc' Expr -> m (String, [String])
forall (m :: * -> *).
MonadError Error m =>
WithLoc' Expr -> m (String, [String])
subscriptExpr WithLoc' Expr
e
        (String, [String], Bool) -> m (String, [String], Bool)
forall (m :: * -> *) a. Monad m => a -> m a
return (String
x, [String]
indices, Bool
starred)
  let pack :: (String, [String], Bool) -> m FormatTree
pack (String
x, [String]
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
$ String -> [String] -> FormatTree
packSubscriptedVar' String
x [String]
indices
        | Bool
otherwise = do
          let xs :: FormatExpr
xs = String -> [String] -> FormatExpr
packSubscriptedVar String
x [String]
indices
          String
i <- VarName -> String
formatVarName (VarName -> String)
-> (WithLoc' VarName -> VarName) -> WithLoc' VarName -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WithLoc' VarName -> VarName
forall a. WithLoc' a -> a
value' (WithLoc' VarName -> String) -> m (WithLoc' VarName) -> m String
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
$ String -> FormatExpr -> FormatTree -> FormatTree
Loop String
i (FormatExpr -> FormatExpr
Len FormatExpr
xs) (String -> [String] -> FormatTree
packSubscriptedVar' String
x ([String]
indices [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ [String
i]))
  case WithLoc' Expr
e of
    CallBuiltin (BuiltinPrint [Type]
_) [WithLoc' Expr]
args -> do
      [(String, [String], Bool)]
outputs <- (WithLoc' Expr -> m (String, [String], Bool))
-> [WithLoc' Expr] -> m [(String, [String], Bool)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM WithLoc' Expr -> m (String, [String], Bool)
forall (m :: * -> *).
MonadError Error m =>
WithLoc' Expr -> m (String, [String], Bool)
starredExpr [WithLoc' Expr]
args
      [FormatTree]
outputs <- ((String, [String], Bool) -> m FormatTree)
-> [(String, [String], Bool)] -> m [FormatTree]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String, [String], Bool) -> m FormatTree
forall (m :: * -> *).
MonadAlpha m =>
(String, [String], Bool) -> m FormatTree
pack [(String, [String], 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 -> String -> m FormatTree
forall (m :: * -> *) a.
MonadError Error m =>
Maybe Loc -> String -> m a
throwSemanticErrorAt' (WithLoc' Expr -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' WithLoc' Expr
e) String
"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 ([String], Either String [String]), FormatTree)
-> m IOFormat
forall (m :: * -> *).
MonadError Error m =>
(FormatTree, Maybe ([String], Either String [String]), FormatTree)
-> m IOFormat
pack ((FormatTree, Maybe ([String], Either String [String]), FormatTree)
 -> m IOFormat)
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
-> m IOFormat
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
go [] [Statement]
body
  where
    pack :: MonadError Error m => (FormatTree, Maybe ([String], Either String [String]), FormatTree) -> m IOFormat
    pack :: (FormatTree, Maybe ([String], Either String [String]), FormatTree)
-> m IOFormat
pack (FormatTree
_, Maybe ([String], Either String [String])
Nothing, FormatTree
_) = String -> m IOFormat
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwSemanticError String
"main function must call solve function"
    pack (FormatTree
inputTree, Just ([String]
inputVariables, Either String [String]
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 :: [String]
-> FormatTree -> Either String [String] -> FormatTree -> IOFormat
IOFormat
          { inputTree :: FormatTree
inputTree = FormatTree
inputTree,
            inputVariables :: [String]
inputVariables = [String]
inputVariables,
            outputVariables :: Either String [String]
outputVariables = Either String [String]
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 ([String], Either String [String]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
go [(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
formats = \case
      Return WithLoc' Expr
_ : [Statement]
_ -> String
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwSemanticError String
"return statement is not allowd in main function"
      AugAssign Target'
_ Operator
_ WithLoc' Expr
_ : [Statement]
_ -> String
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwSemanticError String
"augumented assignment statement is not allowd in main function"
      AnnAssign Target'
x Type
t WithLoc' Expr
e : [Statement]
cont -> do
        (FormatTree
inputs, Maybe ([String], Either String [String])
solve, [Statement]
cont) <- Target'
-> Type
-> WithLoc' Expr
-> [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      [Statement])
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
Target'
-> Type
-> WithLoc' Expr
-> [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      [Statement])
parseAnnAssign Target'
x Type
t WithLoc' Expr
e [Statement]
cont
        [(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
go ([(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
formats [(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
-> [(FormatTree, Maybe ([String], Either String [String]),
     FormatTree)]
-> [(FormatTree, Maybe ([String], Either String [String]),
     FormatTree)]
forall a. [a] -> [a] -> [a]
++ [(FormatTree
inputs, Maybe ([String], Either String [String])
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 ([String], Either String [String]),
       FormatTree))
-> Target'
-> WithLoc' Expr
-> [Statement]
-> m (FormatTree, FormatTree)
forall (m :: * -> *).
MonadError Error m =>
([Statement]
 -> m (FormatTree, Maybe ([String], Either String [String]),
       FormatTree))
-> Target'
-> WithLoc' Expr
-> [Statement]
-> m (FormatTree, FormatTree)
parseFor ([(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
go []) Target'
x WithLoc' Expr
e [Statement]
body
        [(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
go ([(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
formats [(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
-> [(FormatTree, Maybe ([String], Either String [String]),
     FormatTree)]
-> [(FormatTree, Maybe ([String], Either String [String]),
     FormatTree)]
forall a. [a] -> [a] -> [a]
++ [(FormatTree
inputs, Maybe ([String], Either String [String])
forall a. Maybe a
Nothing, FormatTree
outputs)]) [Statement]
cont
      If WithLoc' Expr
_ [Statement]
_ [Statement]
_ : [Statement]
_ -> String
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwSemanticError String
"if statement is not allowd in main function"
      Assert WithLoc' Expr
_ : [Statement]
_ -> String
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwSemanticError String
"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 ([String], Either String [String]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
-> [Statement]
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
go ([(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
formats [(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
-> [(FormatTree, Maybe ([String], Either String [String]),
     FormatTree)]
-> [(FormatTree, Maybe ([String], Either String [String]),
     FormatTree)]
forall a. [a] -> [a] -> [a]
++ [([FormatTree] -> FormatTree
Seq [], Maybe ([String], Either String [String])
forall a. Maybe a
Nothing, FormatTree
output)]) [Statement]
cont
      [] -> do
        let input :: FormatTree
input = [FormatTree] -> FormatTree
Seq (((FormatTree, Maybe ([String], Either String [String]), FormatTree)
 -> FormatTree)
-> [(FormatTree, Maybe ([String], Either String [String]),
     FormatTree)]
-> [FormatTree]
forall a b. (a -> b) -> [a] -> [b]
map (\(FormatTree
x, Maybe ([String], Either String [String])
_, FormatTree
_) -> FormatTree
x) [(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
formats)
        let outputs :: FormatTree
outputs = [FormatTree] -> FormatTree
Seq (((FormatTree, Maybe ([String], Either String [String]), FormatTree)
 -> FormatTree)
-> [(FormatTree, Maybe ([String], Either String [String]),
     FormatTree)]
-> [FormatTree]
forall a b. (a -> b) -> [a] -> [b]
map (\(FormatTree
_, Maybe ([String], Either String [String])
_, FormatTree
z) -> FormatTree
z) [(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
formats)
        Maybe ([String], Either String [String])
solve <- case ((FormatTree, Maybe ([String], Either String [String]), FormatTree)
 -> Maybe ([String], Either String [String]))
-> [(FormatTree, Maybe ([String], Either String [String]),
     FormatTree)]
-> [([String], Either String [String])]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(FormatTree
_, Maybe ([String], Either String [String])
y, FormatTree
_) -> Maybe ([String], Either String [String])
y) [(FormatTree, Maybe ([String], Either String [String]),
  FormatTree)]
formats of
          [] -> Maybe ([String], Either String [String])
-> m (Maybe ([String], Either String [String]))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ([String], Either String [String])
forall a. Maybe a
Nothing
          [([String], Either String [String])
solve] -> Maybe ([String], Either String [String])
-> m (Maybe ([String], Either String [String]))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ([String], Either String [String])
 -> m (Maybe ([String], Either String [String])))
-> Maybe ([String], Either String [String])
-> m (Maybe ([String], Either String [String]))
forall a b. (a -> b) -> a -> b
$ ([String], Either String [String])
-> Maybe ([String], Either String [String])
forall a. a -> Maybe a
Just ([String], Either String [String])
solve
          [([String], Either String [String])]
_ -> String -> m (Maybe ([String], Either String [String]))
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwSemanticError String
"cannot call solve function twice"
        (FormatTree, Maybe ([String], Either String [String]), FormatTree)
-> m (FormatTree, Maybe ([String], Either String [String]),
      FormatTree)
forall (m :: * -> *) a. Monad m => a -> m a
return (FormatTree
input, Maybe ([String], Either String [String])
solve, FormatTree
outputs)

run :: (MonadAlpha m, MonadError Error m) => Program -> m (Maybe IOFormat, Program)
run :: Program -> m (Maybe IOFormat, Program)
run Program
prog = String
-> m (Maybe IOFormat, Program) -> m (Maybe IOFormat, Program)
forall (m :: * -> *) a. MonadError Error m => String -> m a -> m a
wrapError' String
"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)