ParseMain.hs

{-# 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 = \case
  [] -> (Nothing, [])
  ToplevelFunctionDef (WithLoc' loc (VarName (Just "main") Nothing Nothing)) args ret body : stmts -> (Just (loc, args, ret, body), stmts)
  stmt : stmts -> second (stmt :) $ splitMain stmts

checkMainType :: MonadError Error m => MainFunction -> m ()
checkMainType (loc, args, ret, _) = wrapAt' loc $ case args of
  _ : _ -> throwTypeError "main function must not take arguments"
  [] -> case ret of
    VarTy _ -> return ()
    NoneTy -> return ()
    _ -> throwTypeError "main function must return None"

pattern CallBuiltin b args <- WithLoc' _ (Call (WithLoc' _ (Constant (ConstBuiltin b))) args)

pattern CallMethod e a args <- WithLoc' _ (Call (WithLoc' _ (Attribute e a)) args)

pattern IntInput <-
  CallBuiltin (BuiltinInt _) [CallBuiltin BuiltinInput []]

pattern MapIntInputSplit <-
  CallBuiltin
    (BuiltinMap [_] _)
    [ WithLoc' _ (Constant (ConstBuiltin (BuiltinInt _))),
      CallMethod
        (CallBuiltin BuiltinInput [])
        (WithLoc' _ BuiltinSplit)
        []
      ]

pattern ListMapIntInputSplit <-
  CallBuiltin
    (BuiltinList _)
    [ CallBuiltin
        (BuiltinMap [_] _)
        [ WithLoc' _ (Constant (ConstBuiltin (BuiltinInt _))),
          CallMethod
            (CallBuiltin BuiltinInput [])
            (WithLoc' _ BuiltinSplit)
            []
          ]
      ]

pattern 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 x _ e cont = do
  let subscriptTrg x = case value' x of
        NameTrg x -> return (formatVarName (value' x), [])
        SubscriptTrg x (WithLoc' _ (Name i)) -> second (++ [formatVarName (value' i)]) <$> subscriptTrg x
        _ -> throwSemanticErrorAt' (loc' x) $ "name target or subscript target is expected, but got: " ++ formatTarget x
  let subscriptTupleTrg x = case value' x of
        TupleTrg xs -> mapM subscriptTrg xs
        _ -> throwSemanticErrorAt' (loc' x) $ "tuple target is expected, but got: " ++ formatTarget x
  let nameTrg x = case value' x of
        NameTrg x -> return $ formatVarName (value' x)
        _ -> throwSemanticErrorAt' (loc' x) $ "name target is expected, but got: " ++ formatTarget x
  let nameOrTupleTrg x = case value' x of
        NameTrg x -> return . Left $ formatVarName (value' x)
        TupleTrg xs -> Right <$> mapM nameTrg xs
        _ -> throwSemanticErrorAt' (loc' x) $ "name target or tuple target is expected, but got: " ++ formatTarget x
  let nameExpr e = case value' e of
        Name x -> return $ formatVarName (value' x)
        _ -> throwSemanticErrorAt' (loc' e) $ "variable is expected, but got: " ++ formatExpr e
  case e of
    -- int(input())
    IntInput -> do
      (x, indices) <- subscriptTrg x
      return (Seq [packSubscriptedVar' x indices, Newline], Nothing, cont)
    -- map(int, input().split())
    MapIntInputSplit -> do
      outputs <- subscriptTupleTrg x
      return (Seq (map (uncurry packSubscriptedVar') outputs ++ [Newline]), Nothing, cont)
    -- list(map(int, input().split()))
    ListMapIntInputSplit -> do
      (x, indices) <- subscriptTrg x
      case cont of
        Assert (WithLoc' _ (Compare (CallBuiltin (BuiltinLen _) [WithLoc' _ (Name x')]) (CmpOp' Eq' _) n)) : cont | formatVarName (value' x') == x -> do
          i <- formatVarName . value' <$> genVarName'
          n <- nameExpr n
          return (Seq [Loop i (Var n) (Exp (At (packSubscriptedVar x indices) i)), Newline], Nothing, cont)
        _ -> throwSemanticErrorAt' (loc' e) "after `xs = list(map(int, input().split()))', we need to write `assert len(xs) == n`"
    -- list(range(n))
    ListRange n -> do
      let isListRange = \case
            AnnAssign _ _ (ListRange n') | n' == n -> True
            _ -> False
      cont <- return $ dropWhile isListRange cont
      case cont of
        For _ (CallBuiltin BuiltinRange1 [WithLoc' _ (Name n')]) _ : _ | value' n' == n -> return (Seq [], Nothing, cont) -- TODO: add more strict checks
        _ -> throwSemanticErrorAt' (loc' e) "after some repetition of `xs = list(range(n))', we need to write `for i in range(n):`"
    -- solve(...)
    WithLoc' _ (Call (WithLoc' _ (Name (WithLoc' _ (VarName (Just "solve") Nothing Nothing)))) args) -> do
      inputs <- mapM nameExpr args
      output <- nameOrTupleTrg x
      return (Seq [], Just (inputs, output), cont)
    _ -> throwSemanticErrorAt' (loc' e) "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 go x e body = do
  x <- case value' x of
    NameTrg x -> return x
    _ -> throwSemanticErrorAt' (loc' x) $ "for loops in main function must use `range' like `for i in range(n): ...'" ++ formatTarget x
  n <- case e of
    CallBuiltin BuiltinRange1 [n] -> return n
    _ -> throwSemanticErrorAt' (loc' e) $ "for loops in main function must use `range' like `for i in range(n): ...': " ++ formatExpr e
  n <- case value' n of
    Name n -> return $ Right (n, 0)
    BinOp (WithLoc' _ (Name n)) Add (WithLoc' _ (Constant (ConstInt k))) -> return $ Right (n, k)
    BinOp (WithLoc' _ (Name n)) Sub (WithLoc' _ (Constant (ConstInt k))) -> return $ Right (n, - k)
    Call (WithLoc' _ (Constant (ConstBuiltin (BuiltinLen _)))) [WithLoc' _ (Name xs)] -> return $ Left xs
    _ -> throwSemanticErrorAt' (loc' n) $ "for loops in main function must use `range(x)', `range(x + k)', `range(x - k)', `range(len(xs))`: " ++ formatExpr n
  n <- return $ case n of
    Right (n, k) ->
      let n' = Var (formatVarName (value' n))
       in if k == 0 then n' else Plus n' k
    Left xs -> Len (Var (formatVarName (value' xs)))
  (input, solve, output) <- go body
  when (isJust solve) $ do
    throwSemanticError "cannot call `solve(...)' in for loop"
  let x' = formatVarName (value' x)
  return (Loop x' n input, Loop x' n output)

parseExprStatement :: (MonadAlpha m, MonadError Error m) => Expr' -> m FormatTree
parseExprStatement e = do
  let subscriptExpr e = case value' e of
        Name x -> return (formatVarName (value' x), [])
        Subscript e (WithLoc' _ (Name i)) -> second (++ [formatVarName (value' i)]) <$> subscriptExpr e
        _ -> throwSemanticErrorAt' (loc' e) $ "subscripted variable is expected, but got: " ++ formatExpr e
  let starredExpr e = do
        (e, starred) <- return $ case value' e of
          Starred e -> (e, True)
          _ -> (e, False)
        (x, indices) <- subscriptExpr e
        return (x, indices, starred)
  let pack (x, indices, starred)
        | not starred = return $ packSubscriptedVar' x indices
        | otherwise = do
          let xs = packSubscriptedVar x indices
          i <- formatVarName . value' <$> genVarName'
          return $ Loop i (Len xs) (packSubscriptedVar' x (indices ++ [i]))
  case e of
    CallBuiltin (BuiltinPrint _) args -> do
      outputs <- mapM starredExpr args
      outputs <- mapM pack outputs
      return $ Seq (outputs ++ [Newline])
    _ -> throwSemanticErrorAt' (loc' e) "only `print(...)' is allowed for expr statements in main function"

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

run :: (MonadAlpha m, MonadError Error m) => Program -> m (Maybe IOFormat, Program)
run prog = wrapError' "Jikka.RestrictedPython.Convert.ParseMain" $ do
  (main, prog) <- return $ splitMain prog
  main <- forM main $ \main -> do
    checkMainType main
    main <- parseMain main
    return $ normalizeIOFormat main
  return (main, prog)