{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE GADTs #-}
module Analysis.Identity (
  identityAnalysis,
) where

import AST
import Data
import Util.IdGen


-- | Every array, scalar and accumulator has an ID. Trivial values such as
-- Nothing only have the knowledge that they are indeed Nothing. Compound
-- values know which values they consist of.
data ValId t where
  VINil :: ValId TNil
  VIPair :: ValId a -> ValId b -> ValId (TPair a b)
  VIEither :: Either (ValId a) (ValId b) -> ValId (TEither a b)  -- ^ known alternative
  VIEither' :: ValId a -> ValId b -> ValId (TEither a b)  -- ^ unknown alternative, but known values in each case
  VIMaybe :: Maybe (ValId a) -> ValId (TMaybe a)
  VIArr :: Int -> ValId (TArr n t)
  VIScal :: Int -> ValId (TScal t)
  VIAccum :: Int -> ValId (TAccum t)

  -- | We don't know what this consists of, but it's a value, and let's just
  -- give it an ID nevertheless.
  VIThing :: Int -> ValId t

instance Eq (ValId t) where
  VINil         == VINil           = True
  VINil         == _               = False
  VIPair a b    == VIPair a' b'    = a == a' && b == b'
  VIPair{}      == _               = False
  VIEither a    == VIEither a'     = a == a'
  VIEither{}    == _               = False
  VIEither' a b == VIEither' a' b' = a == a' && b == b'
  VIEither'{}   == _               = False
  VIMaybe a     == VIMaybe a'      = a == a'
  VIMaybe{}     == _               = False
  VIArr i       == VIArr i'        = i == i'
  VIArr{}       == _               = False
  VIScal i      == VIScal i'       = i == i'
  VIScal{}      == _               = False
  VIAccum i     == VIAccum i'      = i == i'
  VIAccum{}     == _               = False
  VIThing i     == VIThing i'      = i == i'
  VIThing{}     == _               = False

-- | Symbolic partial evaluation.
identityAnalysis :: SList STy env -> Expr x env t -> Expr ValId env t
identityAnalysis env term = runIdGen 0 $ do
  env' <- slistMapA numberConstant env
  snd <$> idana env' term
  where
    numberConstant :: STy t -> IdGen (ValId t)
    numberConstant = \case
      STNil -> pure VINil
      STPair a b -> VIPair <$> numberConstant a <*> numberConstant b
      STEither a b -> VIEither' <$> numberConstant a <*> numberConstant b
      STMaybe{} -> VIThing <$> genId
      STArr{} -> VIArr <$> genId
      STScal{} -> VIScal <$> genId
      STAccum{} -> VIAccum <$> genId

idana :: SList ValId env -> Expr x env t -> IdGen (ValId t, Expr ValId env t)
idana env = \case
  EVar _ t i -> do
    let v = slistIdx env i
    pure (v, EVar v t i)

  ELet _ e1 e2 -> do
    (v1, e1') <- idana env e1
    (v2, e2') <- idana (v1 `SCons` env) e2
    pure (v2, ELet v2 e1' e2')

  EPair _ e1 e2 -> do
    (v1, e1') <- idana env e1
    (v2, e2') <- idana env e2
    pure (VIPair v1 v2, EPair (VIPair v1 v2) e1' e2')

  EFst _ e -> do
    (v, e') <- idana env e
    v' <- case v of VIPair v1 _ -> pure v1
                    _ -> VIThing <$> genId
    pure (v', EFst v' e')

  ESnd _ e -> do
    (v, e') <- idana env e
    v' <- case v of VIPair _ v2 -> pure v2
                    _ -> VIThing <$> genId
    pure (v', ESnd v' e')

  ENil _ -> pure (VINil, ENil VINil)

  EInl _ t2 e1 -> do
    (v1, e1') <- idana env e1
    let v = VIEither (Left v1)
    pure (v, EInl v t2 e1')

  EInr _ t1 e2 -> do
    (v2, e2') <- idana env e2
    let v = VIEither (Right v2)
    pure (v, EInr v t1 e2')

  ECase _ e1 e2 e3 -> do
    (v1, e1') <- idana env e1
    case v1 of
      VIEither (Left v1') -> do
        (v2, e2') <- idana (v1' `SCons` env) e2
        scrap <- VIThing <$> genId
        (_, e3') <- idana (scrap `SCons` env) e3
        pure (v2, ECase v2 e1' e2' e3')
      VIEither (Right v1') -> do
        scrap <- VIThing <$> genId
        (_, e2') <- idana (scrap `SCons` env) e2
        (v3, e3') <- idana (v1' `SCons` env) e3
        pure (v3, ECase v3 e1' e2' e3')
      VIEither' v1'l v1'r -> do
        (_, e2') <- idana (v1'l `SCons` env) e2
        (_, e3') <- idana (v1'r `SCons` env) e3
        res <- genId
        pure (VIThing res, ECase (VIThing res) e1' e2' e3')
      VIThing _ -> do
        x2 <- genId
        x3 <- genId
        (v2, e2') <- idana (VIThing x2 `SCons` env) e2
        (v3, e3') <- idana (VIThing x3 `SCons` env) e3
        if v2 == v3
          then pure (v2, ECase v2 e1' e2' e3')
          else do
            res <- genId
            pure (VIThing res, ECase (VIThing res) e1' e2' e3')

  ENothing _ t -> pure (VIMaybe Nothing, ENothing (VIMaybe Nothing) t)

  EJust _ e1 -> do
    (v1, e1') <- idana env e1
    let v = VIMaybe (Just v1)
    pure (v, EJust v e1')

  EMaybe _ e1 e2 e3 -> do
    (v3, e3') <- idana env e3
    case v3 of
      VIMaybe Nothing -> do
        (v1, e1') <- idana env e1
        scrap <- VIThing <$> genId
        (_, e2') <- idana (scrap `SCons` env) e2
        pure (v1, EMaybe v1 e1' e2' e3')
      VIMaybe (Just v3j) -> do
        (v2, e2') <- idana (v3j `SCons` env) e2
        (_, e1') <- idana env e1
        pure (v2, EMaybe v2 e1' e2' e3')
      VIThing _ -> do
        (v1, e1') <- idana env e1
        scrap <- VIThing <$> genId
        (v2, e2') <- idana (scrap `SCons` env) e2
        if v1 == v2
          then pure (v2, EMaybe v2 e1' e2' e3')
          else do
            res <- genId
            pure (VIThing res, EMaybe (VIThing res) e1' e2' e3')

  EConstArr _ dim t arr -> do
    x1 <- VIArr <$> genId
    pure (x1, EConstArr x1 dim t arr)

  EBuild _ dim e1 e2 -> do
    (_, e1') <- idana env e1
    scrap <- VIThing <$> genId
    (_, e2') <- idana (scrap `SCons` env) e2
    res <- VIArr <$> genId
    pure (res, EBuild res dim e1' e2')

  EFold1Inner _ e1 e2 e3 -> do
    scrap1 <- VIThing <$> genId
    scrap2 <- VIThing <$> genId
    (_, e1') <- idana (scrap1 `SCons` scrap2 `SCons` env) e1
    (_, e2') <- idana env e2
    (_, e3') <- idana env e3
    res <- VIArr <$> genId
    pure (res, EFold1Inner res e1' e2' e3')

  ESum1Inner _ e1 -> do
    (_, e1') <- idana env e1
    res <- VIArr <$> genId
    pure (res, ESum1Inner res e1')

  EUnit _ e1 -> do
    (_, e1') <- idana env e1
    res <- VIArr <$> genId
    pure (res, EUnit res e1')

  EReplicate1Inner _ e1 e2 -> do
    (_, e1') <- idana env e1
    (_, e2') <- idana env e2
    res <- VIArr <$> genId
    pure (res, EReplicate1Inner res e1' e2')

  EMaximum1Inner _ e1 -> do
    (_, e1') <- idana env e1
    res <- VIArr <$> genId
    pure (res, EMaximum1Inner res e1')

  EMinimum1Inner _ e1 -> do
    (_, e1') <- idana env e1
    res <- VIArr <$> genId
    pure (res, EMinimum1Inner res e1')

  EConst _ t val -> do
    res <- VIScal <$> genId
    pure (res, EConst res t val)

  EIdx0 _ e1 -> do
    (_, e1') <- idana env e1
    res <- VIThing <$> genId
    pure (res, EIdx0 res e1')

  EIdx1 _ e1 e2 -> do
    (_, e1') <- idana env e1
    (_, e2') <- idana env e2
    res <- VIThing <$> genId
    pure (res, EIdx1 res e1' e2')

  EIdx _ e1 e2 -> do
    (_, e1') <- idana env e1
    (_, e2') <- idana env e2
    res <- VIThing <$> genId
    pure (res, EIdx res e1' e2')

  EShape _ e1 -> do
    (_, e1') <- idana env e1
    res <- VIThing <$> genId
    pure (res, EShape res e1')

  EOp _ op e1 -> do
    (_, e1') <- idana env e1
    res <- VIThing <$> genId
    pure (res, EOp res op e1')

  ECustom _ t1 t2 t3 e1 e2 e3 e4 e5 -> do
    x1 <- VIThing <$> genId
    x2 <- VIThing <$> genId
    (_, e1') <- idana (x1 `SCons` x2 `SCons` SNil) e1
    x3 <- VIThing <$> genId
    x4 <- VIThing <$> genId
    (_, e2') <- idana (x3 `SCons` x4 `SCons` SNil) e2
    x5 <- VIThing <$> genId
    x6 <- VIThing <$> genId
    (_, e3') <- idana (x5 `SCons` x6 `SCons` SNil) e3
    (_, e4') <- idana env e4
    (_, e5') <- idana env e5
    res <- VIThing <$> genId
    pure (res, ECustom res t1 t2 t3 e1' e2' e3' e4' e5')

  -- EWith e1 e2

  _ -> _