content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE LambdaCase #-}
  
  module Jikka.RestrictedPython.Language.Util
    ( -- * generating symbols
      genType,
      genVarName,
      genVarName',
  
      -- * free variables
      freeTyVars,
      freeVars,
      freeVars',
      freeVarsTarget,
      freeVarsTarget',
  
      -- * return-statements
      doesAlwaysReturn,
      doesPossiblyReturn,
  
      -- * traversing statements
      mapStatement,
      mapStatementM,
      mapLargeStatement,
      mapLargeStatementM,
      mapStatements,
      mapStatementsM,
      listStatements,
  
      -- * traversing sub exprs
      mapSubExprM,
      mapSubExpr,
      listSubExprs,
  
      -- * traversing exprs
      mapExprTargetM,
      mapExprStatementM,
      mapExprM,
      listExprs,
  
      -- * exprs
      hasFunctionCall,
      isSmallExpr,
      dropLocation,
  
      -- * targets
      targetVars,
      targetVars',
      hasSubscriptTrg,
      hasBareNameTrg,
      exprToTarget,
      targetToExpr,
  
      -- * programs
      toplevelMainDef,
  
      -- * variable names
      mainVarName,
      solveVarName,
      underscoreVarName,
    )
  where
  
  import Control.Monad.Identity
  import Control.Monad.Writer.Strict
  import Data.List (delete, nub)
  import Jikka.Common.Alpha
  import Jikka.Common.Location
  import Jikka.RestrictedPython.Language.Expr
  
  genType :: MonadAlpha m => m Type
  genType = VarTy . TypeName Nothing . Just <$> nextCounter
  
  genVarName :: MonadAlpha m => VarName' -> m VarName'
  genVarName x@(WithLoc' _ (VarName occ _ hint)) = do
    i <- nextCounter
    return . WithLoc' (loc' x) $ VarName occ (Just i) hint
  
  genVarName' :: MonadAlpha m => m VarName'
  genVarName' = genVarName (withoutLoc underscoreVarName)
  
  freeTyVars :: Type -> [TypeName]
  freeTyVars = nub . go
    where
      go = \case
        VarTy x -> [x]
        IntTy -> []
        BoolTy -> []
        ListTy t -> go t
        TupleTy ts -> concat $ mapM go ts
        CallableTy ts ret -> concat $ mapM go (ret : ts)
        StringTy -> []
        SideEffectTy -> []
  
  -- | `freeVars'` reports all free variables.
  freeVars :: Expr' -> [VarName]
  freeVars = nub . map value' . freeVars'
  
  -- | `freeVars'` reports all free variables with their locations, i.e. occurrences.
  -- For examples, @x + x@ and @x@ have the same free variables @x@ but they have different sets of occurrences of free variable.
  freeVars' :: Expr' -> [VarName']
  freeVars' (WithLoc' _ e0) = case e0 of
    BoolOp e1 _ e2 -> freeVars' e1 ++ freeVars' e2
    BinOp e1 _ e2 -> freeVars' e1 ++ freeVars' e2
    UnaryOp _ e -> freeVars' e
    Lambda args e -> foldl (\vars (x, _) -> delete x vars) (freeVars' e) args
    IfExp e1 e2 e3 -> freeVars' e1 ++ freeVars' e2 ++ freeVars' e3
    ListComp e (Comprehension x iter pred) -> freeVars' iter ++ foldl (\vars x -> delete x vars) (freeVars' e ++ concatMap freeVars' pred) (targetVars' x)
    Compare e1 _ e2 -> freeVars' e1 ++ freeVars' e2
    Call f args -> concatMap freeVars' (f : args)
    Constant _ -> []
    Attribute e _ -> freeVars' e
    Subscript e1 e2 -> freeVars' e1 ++ freeVars' e2
    Starred e -> freeVars' e
    Name x -> [x]
    List _ es -> concatMap freeVars' es
    Tuple es -> concatMap freeVars' es
    SubscriptSlice e from to step -> freeVars' e ++ concatMap freeVars' from ++ concatMap freeVars' to ++ concatMap freeVars' step
  
  freeVarsTarget :: Target' -> [VarName]
  freeVarsTarget = nub . map value' . freeVarsTarget'
  
  freeVarsTarget' :: Target' -> [VarName']
  freeVarsTarget' (WithLoc' _ x) = case x of
    SubscriptTrg _ e -> freeVars' e
    NameTrg _ -> []
    TupleTrg xs -> concatMap freeVarsTarget' xs
  
  doesAlwaysReturn :: Statement -> Bool
  doesAlwaysReturn = \case
    Return _ -> True
    AugAssign _ _ _ -> False
    AnnAssign _ _ _ -> False
    For _ _ _ -> False
    If _ body1 body2 -> any doesAlwaysReturn body1 && any doesAlwaysReturn body2
    Assert _ -> False
    Expr' _ -> False
  
  doesPossiblyReturn :: Statement -> Bool
  doesPossiblyReturn = \case
    Return _ -> True
    AugAssign _ _ _ -> False
    AnnAssign _ _ _ -> False
    For _ _ body -> any doesPossiblyReturn body
    If _ body1 body2 -> any doesPossiblyReturn body1 || any doesPossiblyReturn body2
    Assert _ -> False
    Expr' _ -> False
  
  -- | `mapSubExprM` replaces all exprs in a given expr using a given function.
  -- This may breaks various constraints.
  mapSubExprM :: Monad m => (Expr' -> m Expr') -> Expr' -> m Expr'
  mapSubExprM f = go
    where
      go e0 =
        f . WithLoc' (loc' e0) =<< case value' e0 of
          BoolOp e1 op e2 -> BoolOp <$> go e1 <*> return op <*> go e2
          BinOp e1 op e2 -> BinOp <$> go e1 <*> return op <*> go e2
          UnaryOp op e -> UnaryOp op <$> go e
          Lambda args body -> Lambda args <$> go body
          IfExp e1 e2 e3 -> IfExp <$> go e1 <*> go e2 <*> go e3
          ListComp e (Comprehension x iter pred) -> do
            e <- go e
            x <- mapExprTargetM f x
            iter <- go iter
            pred <- mapM go pred
            return $ ListComp e (Comprehension x iter pred)
          Compare e1 op e2 -> Compare <$> go e1 <*> return op <*> go e2
          Call g args -> Call <$> go g <*> mapM go args
          Constant const -> return $ Constant const
          Attribute e x -> Attribute <$> go e <*> pure x
          Subscript e1 e2 -> Subscript <$> go e1 <*> go e2
          Starred e -> Starred <$> go e
          Name x -> return $ Name x
          List t es -> List t <$> mapM go es
          Tuple es -> Tuple <$> mapM go es
          SubscriptSlice e from to step -> SubscriptSlice <$> go e <*> mapM go from <*> mapM go to <*> mapM go step
  
  mapSubExpr :: (Expr' -> Expr') -> Expr' -> Expr'
  mapSubExpr f = runIdentity . mapSubExprM (return . f)
  
  listSubExprs :: Expr' -> [Expr']
  listSubExprs = reverse . getDual . execWriter . mapSubExprM go
    where
      go e = do
        tell $ Dual [e]
        return e
  
  mapExprTargetM :: Monad m => (Expr' -> m Expr') -> Target' -> m Target'
  mapExprTargetM f x =
    WithLoc' (loc' x) <$> case value' x of
      SubscriptTrg x e -> SubscriptTrg <$> mapExprTargetM f x <*> f e
      NameTrg x -> return $ NameTrg x
      TupleTrg xs -> TupleTrg <$> mapM (mapExprTargetM f) xs
  
  mapExprStatementM :: Monad m => (Expr' -> m Expr') -> Statement -> m Statement
  mapExprStatementM f = \case
    Return e -> Return <$> f e
    AugAssign x op e -> AugAssign <$> mapExprTargetM f x <*> pure op <*> f e
    AnnAssign x t e -> AnnAssign <$> mapExprTargetM f x <*> pure t <*> f e
    For x iter body -> For <$> mapExprTargetM f x <*> f iter <*> mapM (mapExprStatementM f) body
    If e body1 body2 -> If <$> f e <*> mapM (mapExprStatementM f) body1 <*> mapM (mapExprStatementM f) body2
    Assert e -> Assert <$> f e
    Expr' e -> Expr' <$> f e
  
  mapExprToplevelStatementM :: Monad m => (Expr' -> m Expr') -> ToplevelStatement -> m ToplevelStatement
  mapExprToplevelStatementM f = \case
    ToplevelAnnAssign x t e -> ToplevelAnnAssign x t <$> f e
    ToplevelFunctionDef g args ret body -> ToplevelFunctionDef g args ret <$> mapM (mapExprStatementM f) body
    ToplevelAssert e -> ToplevelAssert <$> f e
  
  mapExprM :: Monad m => (Expr' -> m Expr') -> Program -> m Program
  mapExprM f = mapM (mapExprToplevelStatementM f)
  
  listExprs :: Program -> [Expr']
  listExprs = reverse . getDual . execWriter . mapExprM go
    where
      go e = do
        tell $ Dual [e]
        return e
  
  mapStatementStatementM :: Monad m => (Statement -> m [Statement]) -> Statement -> m [Statement]
  mapStatementStatementM f = \case
    Return e -> f $ Return e
    AugAssign x op e -> f $ AugAssign x op e
    AnnAssign x t e -> f $ AnnAssign x t e
    For x iter body -> do
      body <- concat <$> mapM (mapStatementStatementM f) body
      f $ For x iter body
    If e body1 body2 -> do
      body1 <- concat <$> mapM (mapStatementStatementM f) body1
      body2 <- concat <$> mapM (mapStatementStatementM f) body2
      f $ If e body1 body2
    Assert e -> f $ Assert e
    Expr' e -> f $ Expr' e
  
  mapStatementToplevelStatementM :: Monad m => (Statement -> m [Statement]) -> ToplevelStatement -> m ToplevelStatement
  mapStatementToplevelStatementM go = \case
    ToplevelAnnAssign x t e -> return $ ToplevelAnnAssign x t e
    ToplevelFunctionDef f args ret body -> do
      body <- concat <$> mapM (mapStatementStatementM go) body
      return $ ToplevelFunctionDef f args ret body
    ToplevelAssert e -> return $ ToplevelAssert e
  
  -- | `mapStatementM` replaces all statements in a given program using a given function.
  -- This may breaks various constraints.
  mapStatementM :: Monad m => (Statement -> m [Statement]) -> Program -> m Program
  mapStatementM f = mapM (mapStatementToplevelStatementM f)
  
  mapStatement :: (Statement -> [Statement]) -> Program -> Program
  mapStatement f = runIdentity . mapStatementM (return . f)
  
  mapLargeStatementM :: Monad m => (Expr' -> [Statement] -> [Statement] -> m [Statement]) -> (Target' -> Expr' -> [Statement] -> m [Statement]) -> Program -> m Program
  mapLargeStatementM fIf fFor = mapStatementM go
    where
      go = \case
        Return e -> return [Return e]
        AugAssign x op e -> return [AugAssign x op e]
        AnnAssign x t e -> return [AnnAssign x t e]
        For x iter body -> fFor x iter body
        If e body1 body2 -> fIf e body1 body2
        Assert e -> return [Assert e]
        Expr' e -> return [Expr' e]
  
  mapLargeStatement :: (Expr' -> [Statement] -> [Statement] -> [Statement]) -> (Target' -> Expr' -> [Statement] -> [Statement]) -> Program -> Program
  mapLargeStatement fIf fFor = runIdentity . mapLargeStatementM fIf' fFor'
    where
      fIf' e body1 body2 = return $ fIf e body1 body2
      fFor' x iter body = return $ fFor x iter body
  
  listStatements :: Program -> [Statement]
  listStatements = reverse . getDual . execWriter . mapStatementM go
    where
      go stmt = do
        tell $ Dual [stmt]
        return [stmt]
  
  mapStatementsToplevelStatementM :: Monad m => ([Statement] -> m [Statement]) -> ToplevelStatement -> m ToplevelStatement
  mapStatementsToplevelStatementM go = \case
    ToplevelAnnAssign x t e -> return $ ToplevelAnnAssign x t e
    ToplevelFunctionDef f args ret body -> do
      let go' = \case
            Return e -> return [Return e]
            AugAssign x op e -> return [AugAssign x op e]
            AnnAssign x t e -> return [AnnAssign x t e]
            For x iter body -> do
              body <- go body
              return [For x iter body]
            If e body1 body2 -> do
              body1 <- go body1
              body2 <- go body2
              return [If e body1 body2]
            Assert e -> return [Assert e]
            Expr' e -> return [Expr' e]
      body <- concat <$> mapM (mapStatementStatementM go') body
      body <- go body
      return $ ToplevelFunctionDef f args ret body
    ToplevelAssert e -> return $ ToplevelAssert e
  
  mapStatementsM :: Monad m => ([Statement] -> m [Statement]) -> Program -> m Program
  mapStatementsM f = mapM (mapStatementsToplevelStatementM f)
  
  mapStatements :: ([Statement] -> [Statement]) -> Program -> Program
  mapStatements f = runIdentity . mapStatementsM (return . f)
  
  hasFunctionCall :: Expr' -> Bool
  hasFunctionCall = any (check . value') . listSubExprs
    where
      check = \case
        Call _ _ -> True
        _ -> False
  
  -- | `isSmallExpr` is true if the evaluation of a given expr trivially terminates.
  isSmallExpr :: Expr' -> Bool
  isSmallExpr = not . hasFunctionCall
  
  dropLocation :: Expr' -> Expr'
  dropLocation = mapSubExpr go
    where
      go (WithLoc' _ e) = withoutLoc e
  
  targetVars :: Target' -> [VarName]
  targetVars = nub . map value' . targetVars'
  
  targetVars' :: Target' -> [VarName']
  targetVars' (WithLoc' _ x) = case x of
    SubscriptTrg x _ -> targetVars' x
    NameTrg x -> [x]
    TupleTrg xs -> concatMap targetVars' xs
  
  hasSubscriptTrg :: Target' -> Bool
  hasSubscriptTrg (WithLoc' _ x) = case x of
    SubscriptTrg _ _ -> True
    NameTrg _ -> False
    TupleTrg xs -> any hasSubscriptTrg xs
  
  hasBareNameTrg :: Target' -> Bool
  hasBareNameTrg (WithLoc' _ x) = case x of
    SubscriptTrg _ _ -> False
    NameTrg _ -> True
    TupleTrg xs -> any hasSubscriptTrg xs
  
  exprToTarget :: Expr' -> Maybe Target'
  exprToTarget e =
    WithLoc' (loc' e) <$> case value' e of
      Name x -> Just $ NameTrg x
      Tuple es -> TupleTrg <$> mapM exprToTarget es
      Subscript e1 e2 -> SubscriptTrg <$> exprToTarget e1 <*> pure e2
      _ -> Nothing
  
  targetToExpr :: Target' -> Expr'
  targetToExpr e =
    WithLoc' (loc' e) $ case value' e of
      NameTrg x -> Name x
      TupleTrg es -> Tuple (map targetToExpr es)
      SubscriptTrg e1 e2 -> Subscript (targetToExpr e1) e2
  
  toplevelMainDef :: [Statement] -> Program
  toplevelMainDef body = [ToplevelFunctionDef (WithLoc' Nothing mainVarName) [] IntTy body]
  
  mainVarName :: VarName
  mainVarName = VarName (Just "main") Nothing Nothing
  
  solveVarName :: VarName
  solveVarName = VarName (Just "solve") Nothing Nothing
  
  underscoreVarName :: VarName
  underscoreVarName = VarName Nothing Nothing Nothing
  