{-# LANGUAGE LambdaCase #-}

module Jikka.RestrictedPython.Language.VariableAnalysis where

import Data.List (delete, intersect, nub)
import Jikka.RestrictedPython.Language.Expr
import Jikka.RestrictedPython.Language.Util

newtype ReadList = ReadList [VarName]
  deriving (Eq, Ord, Show, Read)

newtype WriteList = WriteList [VarName]
  deriving (Eq, Ord, Show, Read)

haveWriteReadIntersection :: WriteList -> ReadList -> Bool
haveWriteReadIntersection (WriteList w) (ReadList r) = not (null (w `intersect` r))

analyzeExpr :: Expr' -> ReadList
analyzeExpr = ReadList . freeVars

analyzeTargetRead :: Target' -> ReadList
analyzeTargetRead = ReadList . freeVarsTarget

analyzeTargetWrite :: Target' -> WriteList
analyzeTargetWrite = WriteList . targetVars

analyzeStatementGeneric :: Bool -> Statement -> (ReadList, WriteList)
analyzeStatementGeneric isMax = \case
  Return e -> (analyzeExpr e, WriteList [])
  AugAssign x _ e ->
    let w = analyzeTargetWrite x
        (ReadList r) = analyzeTargetRead x
        (ReadList r') = analyzeExpr e
     in (ReadList (nub $ r ++ r'), w)
  AnnAssign x _ e ->
    let w = analyzeTargetWrite x
        (ReadList r) = analyzeTargetRead x
        (ReadList r') = analyzeExpr e
     in (ReadList (nub $ r ++ r'), w)
  For x iter body ->
    let xs = targetVars x
        ReadList r = analyzeExpr iter
        (ReadList r', WriteList w) = analyzeStatementsGeneric isMax body
     in if isMax
          then (ReadList (nub $ r ++ foldl (flip delete) r' xs), WriteList (nub $ foldl (flip delete) w xs))
          else (ReadList r, WriteList [])
  If e body1 body2 ->
    let ReadList r = analyzeExpr e
        (ReadList r1, WriteList w1) = analyzeStatementsGeneric isMax body1
        (ReadList r2, WriteList w2) = analyzeStatementsGeneric isMax body2
     in if isMax
          then (ReadList (nub $ r ++ r1 ++ r2), WriteList (nub $ w1 ++ w2))
          else (ReadList (nub $ r ++ intersect r1 r2), WriteList (nub $ w1 `intersect` w2))
  Assert e -> (analyzeExpr e, WriteList [])
  Append _ _ x e ->
    let w = maybe (WriteList []) analyzeTargetWrite (exprToTarget x)
        (ReadList r) = maybe (ReadList []) analyzeTargetRead (exprToTarget x)
        (ReadList r') = analyzeExpr e
     in (ReadList (nub $ r ++ r'), w)
  Expr' e -> (analyzeExpr e, WriteList [])

analyzeStatementsGeneric :: Bool -> [Statement] -> (ReadList, WriteList)
analyzeStatementsGeneric isMax = go [] []
  where
    go r w [] = (ReadList (nub r), WriteList (nub w))
    go r w (stmt : stmts) =
      let (ReadList r', WriteList w') = analyzeStatementGeneric isMax stmt
       in go (r' ++ r) (w' ++ w) stmts

-- | `analyzeStatementMax` returns lists of variables which are possibly read or written in given statements.
analyzeStatementMax :: Statement -> (ReadList, WriteList)
analyzeStatementMax = analyzeStatementGeneric True

analyzeStatementsMax :: [Statement] -> (ReadList, WriteList)
analyzeStatementsMax = analyzeStatementsGeneric True

-- | `analyzeStatementMin` returns lists of variables which are always read or written in given statements.
analyzeStatementMin :: Statement -> (ReadList, WriteList)
analyzeStatementMin = analyzeStatementGeneric False

analyzeStatementsMin :: [Statement] -> (ReadList, WriteList)
analyzeStatementsMin = analyzeStatementsGeneric False