{-# LANGUAGE FlexibleContexts #-}

-- |
-- Module      : Jikka.RestrictedPython.Convert.TypeInfer
-- Description : split a for-loop into many small for-loops based on the dependency graph of variables and assignments. / 変数と代入の依存関係グラフに基づいて、for ループを複数の小さな for ループに分割します。
-- Copyright   : (c) Kimiyuki Onaka, 2021
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
module Jikka.RestrictedPython.Convert.SplitLoops
  ( run,
    run',
    runForLoop,
  )
where

import Data.List (partition)
import Jikka.Common.Alpha
import Jikka.Common.Error
import qualified Jikka.RestrictedPython.Convert.Alpha as Alpha (run)
import Jikka.RestrictedPython.Language.Expr
import Jikka.RestrictedPython.Language.Lint
import Jikka.RestrictedPython.Language.Util
import Jikka.RestrictedPython.Language.VariableAnalysis

-- | `runForLoop` splits a for-loop to many for-loops as possible.
-- This assumes that `doesntHaveSubscriptionInLoopCounters`, `doesntHaveAssignmentToLoopCounters`, and `doesntHaveAssignmentToLoopIterators` hold.
--
-- This function analyzes read-variables and write-variables in statements, and split statements into connected components.
runForLoop :: Target' -> Expr' -> [Statement] -> [Statement]
runForLoop x iter body =
  let connected (_, (r, w)) (_, (r', w')) = haveWriteReadIntersection w r' || haveWriteReadIntersection w' r
      go result [] = reverse result
      go result (stmt : stmts) =
        let (same, diff) = partition (connected stmt) stmts
         in go (For x iter (map fst (stmt : same)) : result) diff
      body' = map (\stmt -> (stmt, analyzeStatementMax stmt)) body
   in go [] body'

-- | `run'` splits for-loops into many small for-loops as possible.
-- This assumes that `doesntHaveSubscriptionInLoopCounters`, `doesntHaveAssignmentToLoopCounters`, and `doesntHaveAssignmentToLoopIterators` hold.
-- This may introduce name conflicts.
--
-- For example, the following
--
-- > a = 0
-- > b = 0
-- > for i in range(10):
-- >     c = b
-- >     a += i
-- >     b += c
--
-- is split to
--
-- > a = 0
-- > b = 0
-- > for i in range(10):
-- >     c = b
-- >     b += c
-- > for i in range(10):
-- >     a += i
run' :: Program -> Program
run' = mapLargeStatement (\e pred1 pred2 -> [If e pred1 pred2]) runForLoop

-- | `run` does alpha conversion, check assumptions, and `run'`.
run :: (MonadAlpha m, MonadError Error m) => Program -> m Program
run prog = wrapError' "Jikka.RestrictedPython.Convert.SplitLoops" $ do
  prog <- Alpha.run prog
  ensureDoesntHaveSubscriptionInLoopCounters prog
  ensureDoesntHaveAssignmentToLoopCounters prog
  ensureDoesntHaveAssignmentToLoopIterators prog
  prog <- return $ run' prog
  Alpha.run prog