module Utils.TH.AlphaConversion
  ( alphaConvE
  ) where

import Import.NoModel
  
import Language.Haskell.TH

import qualified Data.Map.Strict as Map
import qualified Data.Set as Set


alphaConvE :: Map Name Name -> Exp -> Exp
alphaConvE ns = alphaConvE' Set.empty
  where
    alphaConvE' bound = \case
      VarE varn
        | Just varn' <- Map.lookup varn $ ns `Map.withoutKeys` bound
          -> VarE varn'
      o@(VarE _) -> o
      o@(ConE _) -> o
      o@(LitE _) -> o
      AppE f x -> AppE (alphaConvE' bound f) (alphaConvE' bound x)
      AppTypeE f t -> AppTypeE (alphaConvE' bound f) t
      InfixE mA o mB -> InfixE (alphaConvE' bound <$> mA) (alphaConvE' bound o) (alphaConvE' bound <$> mB)
      UInfixE a o b -> UInfixE (alphaConvE' bound a) (alphaConvE' bound o) (alphaConvE' bound b)
      ParensE e -> ParensE $ alphaConvE' bound e
      LamE ps e -> LamE ps $ alphaConvE' (bound `Set.union` foldMap pVars ps) e
      LamCaseE _ -> error "alphaConvE: LamCaseE not implemented"
      TupE es -> TupE $ over (traverse . _Just) (alphaConvE' bound) es
      UnboxedTupE es -> UnboxedTupE $ over (traverse . _Just) (alphaConvE' bound) es
      UnboxedSumE e sAlt sAry -> UnboxedSumE (alphaConvE' bound e) sAlt sAry
      CondE i t e -> CondE (alphaConvE' bound i) (alphaConvE' bound t) (alphaConvE' bound e)
      MultiIfE _ -> error "alphaConvE: MultiIfE not implemented"
      LetE _ _ -> error "alphaConvE: LetE not implemented"
      CaseE _ _ -> error "alphaConvE: CaseE not implemented"
      DoE _ -> error "alphaConvE: DoE not implemented"
      MDoE _ -> error "alphaConvE: MDoE not implemented"
      CompE _ -> error "alphaConvE: CompE not implemented"
      ArithSeqE (FromR e) -> ArithSeqE . FromR $ alphaConvE' bound e
      ArithSeqE (FromThenR e1 e2) -> ArithSeqE $ FromThenR (alphaConvE' bound e1) (alphaConvE' bound e2)
      ArithSeqE (FromToR e1 e2) -> ArithSeqE $ FromToR (alphaConvE' bound e1) (alphaConvE' bound e2)
      ArithSeqE (FromThenToR e1 e2 e3) -> ArithSeqE $ FromThenToR (alphaConvE' bound e1) (alphaConvE' bound e2) (alphaConvE' bound e3)
      ListE es -> ListE $ map (alphaConvE' bound) es
      SigE e t -> SigE (alphaConvE' bound e) t
      RecConE rN es -> RecConE rN $ over (traverse . _2) (alphaConvE' bound) es
      RecUpdE e es -> RecUpdE (alphaConvE' bound e) $ over (traverse . _2) (alphaConvE' bound) es
      StaticE e -> StaticE $ alphaConvE' bound e
      o@(UnboundVarE _) -> o
      o@(LabelE _) -> o
      o@(ImplicitParamVarE _) -> o

pVars :: Pat -> Set Name
pVars = \case
  LitP _            -> Set.empty
  VarP n            -> Set.singleton n
  TupP ps           -> foldMap pVars ps
  UnboxedTupP ps    -> foldMap pVars ps
  UnboxedSumP p _ _ -> pVars p
  ConP _ ps         -> foldMap pVars ps
  InfixP p1 _ p2    -> pVars p1 `Set.union` pVars p2
  UInfixP p1 _ p2   -> pVars p1 `Set.union` pVars p2
  ParensP p         -> pVars p
  TildeP p          -> pVars p
  BangP p           -> pVars p
  AsP n p           -> Set.insert n $ pVars p
  WildP             -> Set.empty
  RecP _ ps         -> foldMapOf (folded . _2) pVars ps
  ListP ps          -> foldMap pVars ps
  SigP p _          -> pVars p
  ViewP _ p         -> pVars p
