content: {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE Strict #-}
  {-# OPTIONS_GHC -fno-warn-name-shadowing #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE LambdaCase #-}
  
  -- | The substitution stepper.
  
  module Duet.Stepper
    ( expandSeq1
    , fargs
    , lookupNameByString
    ) where
  
  import           Control.Applicative
  import           Control.Monad.Catch
  import           Control.Monad.State
  import           Data.List
  import           Data.Map.Strict (Map)
  import qualified Data.Map.Strict as M
  import           Data.Maybe
  import           Duet.Types
  
  --------------------------------------------------------------------------------
  -- Expansion
  
  expandSeq1
    :: (MonadThrow m)
    => Context Type Name (Location)
    -> [BindGroup Type Name (Location)]
    -> Expression Type Name (Location)
    -> m (Expression Type Name (Location))
  expandSeq1 (Context { contextTypeClasses = typeClassEnv
                      , contextSpecialSigs = specialSigs
                      , contextSignatures = signatures
                      }) b e = evalStateT (go e) False
    where
      go =
        \case
          e0
            -- If we're looking at a constructor, then force the args.
            | (ce@(ConstructorExpression l _), args) <- fargs e0 -> do
              args' <- mapM go args
              pure (foldl (ApplicationExpression l) ce args')
            -- If we're looking at a constant (hole), then force the args.
            | (ce@(ConstantExpression l _), args) <- fargs e0 -> do
              args' <- mapM go args
              pure (foldl (ApplicationExpression l) ce args')
            -- We're looking at a general expression, check if a force
            -- has already happpened. If so, we just return the
            -- identity.
            | otherwise -> do
              alreadyExpanded <- get
              if alreadyExpanded
                then pure e0
                else do
                  -- If we haven't expanded anything yet, let's expand
                  -- this mother to whnf.
                  e' <- lift (expandWhnf typeClassEnv specialSigs signatures e0 b)
                  -- If the expansion did actually produce a new AST
                  -- then count that as an expansion.
                  put (e' /= e0)
                  pure e'
  
  expandWhnf
    :: MonadThrow m
    => Map Name (Class Type Name (TypeSignature Type Name Location))
    -> SpecialSigs Name
    -> [TypeSignature Type Name Name]
    -> Expression Type Name (Location)
    -> [BindGroup Type Name (Location)]
    -> m (Expression Type Name (Location))
  expandWhnf typeClassEnv specialSigs signatures e b = go e
    where
      go x =
        case x of
          -- Parens aren't an expansion step, just a grouping.
          ParensExpression _ e -> go e
          VariableExpression _ i -> do
            case find ((== i) . typeSignatureA) signatures of
              Nothing -> do
                e' <- lookupName i b
                pure e'
              Just {} -> pure x
          LiteralExpression {} -> return x
          ConstructorExpression {} -> return x
          ConstantExpression {} -> return x
          ApplicationExpression l (ApplicationExpression l1 op@(VariableExpression _ (PrimopName primop)) x) y ->
            case x of
              LiteralExpression _ (StringLiteral sx) ->
                case y of
                  LiteralExpression _ (StringLiteral sy) ->
                    case primop of
                      PrimopStringAppend ->
                        pure (LiteralExpression l (StringLiteral (sx <> sy)))
                      _ -> error "Runtime type error that should not occurr"
                  _ -> do
                    y' <- go y
                    pure
                      (ApplicationExpression l (ApplicationExpression l1 op x) y')
              LiteralExpression _ (IntegerLiteral n) ->
                case y of
                  LiteralExpression _ (StringLiteral sy) ->
                    case primop of
                      PrimopStringTake ->
                        pure (LiteralExpression l (StringLiteral (genericTake n sy)))
                      PrimopStringDrop ->
                        pure (LiteralExpression l (StringLiteral (genericDrop n sy)))
                      _ -> error "Runtime type error that should not occurr"
                  _ -> do
                    y' <- go y
                    pure
                      (ApplicationExpression l (ApplicationExpression l1 op x) y')
              _ -> do
                x' <- go x
                pure (ApplicationExpression l (ApplicationExpression l1 op x') y)
          ApplicationExpression l func arg ->
            case func of
              LambdaExpression l0 (Alternative l' params body) ->
                case params of
                  (VariablePattern _ param:params') ->
                    let body' = substitute param arg body
                    in case params' of
                         [] -> pure body'
                         _ ->
                           pure
                             (LambdaExpression l0 (Alternative l' params' body'))
                  _ -> error "Unsupported lambda."
              VariableExpression _ (MethodName _ methodName) ->
                case arg of
                  VariableExpression _ dictName@DictName {} ->
                    case find
                           ((== dictName) . dictionaryName)
                           (concatMap
                              (map instanceDictionary . classInstances)
                              (M.elems typeClassEnv)) of
                      Nothing -> throwM (CouldntFindMethodDict dictName)
                      Just dict ->
                        case M.lookup
                               methodName
                               (M.mapKeys
                                  (\(MethodName _ s) -> s)
                                  (dictionaryMethods dict)) of
                          Nothing ->
                            error
                              ("Missing method " ++
                               show methodName ++ " in dictionary: " ++ show dict)
                          Just (_, Alternative _ _ e) -> pure (fmap typeSignatureA e)
                  _ -> error "Unsupported variable expression."
              _ -> do
                func' <- go func
                pure (ApplicationExpression l func' arg)
          orig@(InfixExpression l x op@(_s, VariableExpression _ (PrimopName primop)) y) ->
            case x of
              LiteralExpression _ x' ->
                case y of
                  LiteralExpression _ y' ->
                    case (x', y') of
                      (IntegerLiteral i1, IntegerLiteral i2) ->
                        pure
                          (LiteralExpression
                             l
                             (case primop of
                                PrimopIntegerPlus -> IntegerLiteral (i1 + i2)
                                PrimopIntegerTimes -> IntegerLiteral (i1 * i2)
                                PrimopIntegerSubtract -> IntegerLiteral (i1 - i2)
                                _ -> error "Unexpected operation for integer literals."))
                      (RationalLiteral i1, RationalLiteral i2) ->
                        pure
                          (LiteralExpression
                             l
                             (case primop of
                                PrimopRationalPlus -> RationalLiteral (i1 + i2)
                                PrimopRationalTimes -> RationalLiteral (i1 * i2)
                                PrimopRationalSubtract ->
                                  RationalLiteral (i1 - i2)
                                PrimopRationalDivide -> RationalLiteral (i1 / i2)
                                _ -> error "Unexpected operation for rational literals."))
                      _ -> pure orig
                  _ -> do
                    y' <- go y
                    pure (InfixExpression l x op y')
              _ -> do
                x' <- go x
                pure (InfixExpression l x' op y)
          InfixExpression l x (s, op) y -> do
            op' <- go op
            pure (InfixExpression l x (s, op') y)
          IfExpression l pr th el ->
            case pr of
              ConstructorExpression _ n
                | n == specialSigsTrue specialSigs -> pure th
                | n == specialSigsFalse specialSigs -> pure el
              _ -> IfExpression l <$> go pr <*> pure th <*> pure el
          LetExpression {} -> return x
          LambdaExpression {} -> return x
          CaseExpression l e0 alts ->
            let matches =
                  map
                    (\ca -> (match e0 (caseAltPattern ca), caseAltExpression ca))
                    alts
            in case listToMaybe
                      (mapMaybe
                         (\(r, e) -> do
                            case r of
                              OK v -> pure (v, e)
                              Fail -> Nothing)
                         matches) of
                 Just (Success subs, expr) ->
                   return
                     (foldr
                        (\(name, that) expr' -> substitute name that expr')
                        expr
                        subs)
                 Just (NeedsMoreEval is, _) -> do
                   e' <- expandAt typeClassEnv is specialSigs signatures e0 b
                   pure (CaseExpression l e' alts)
                 Nothing -> error ("Incomplete pattern match... " ++ show matches)
  
  expandAt
    :: MonadThrow m
    => Map Name (Class Type Name (TypeSignature Type Name Location))
    -> [Int]
    -> SpecialSigs Name
    -> [TypeSignature Type Name Name]
    -> Expression Type Name (Location)
    -> [BindGroup Type Name (Location)]
    -> m (Expression Type Name (Location))
  expandAt typeClassEnv is specialSigs signatures e0 b  = go [0] e0
    where
      go js e =
        if is == js
          then expandWhnf typeClassEnv specialSigs signatures e b
          else case e of
                 _
                   | (ce@(ConstructorExpression l _), args) <- fargs e -> do
                     args' <-
                       sequence
                         (zipWith (\i arg -> go (js ++ [i]) arg) [0 ..] args)
                     pure (foldl (ApplicationExpression l) ce args')
                   | otherwise -> pure e
  
  --------------------------------------------------------------------------------
  -- Pattern matching
  
  match
    :: (Eq i)
    => Expression Type i l -> Pattern Type i l -> Result (Match Type i l)
  match = go [0]
    where
      go is val pat =
        case pat of
          BangPattern p
            | isWhnf val -> go is val p
            | otherwise -> OK (NeedsMoreEval is)
          AsPattern _l ident pat ->
            case go is val pat of
              OK (Success binds) -> OK (Success ((ident, val) : binds))
              res -> res
          WildcardPattern _ _ -> OK (Success [])
          VariablePattern _ i -> OK (Success [(i, val)])
          LiteralPattern _ l ->
            case val of
              LiteralExpression _ l'
                | l' == l -> OK (Success [])
                | otherwise -> Fail
              _ -> OK (NeedsMoreEval is)
          ConstructorPattern _ i pats
            | (constructor@ConstructorExpression {}, args) <- fargs val ->
              if fmap (const ()) constructor == ConstructorExpression () i
                then if length args == length pats
                       then foldl
                              (<>)
                              (OK (Success []))
                              (zipWith
                                 (\j (arg, p) -> go (is ++ [j]) arg p)
                                 [0 ..]
                                 (zip args pats))
                       else Fail
                else Fail
            | otherwise -> OK (NeedsMoreEval is)
  
  isWhnf :: Expression Type i l -> Bool
  isWhnf =
    \case
      VariableExpression {} -> False
      ConstructorExpression {} -> True
      ConstantExpression {} -> True
      LiteralExpression {} -> True
      ApplicationExpression {} -> False
      InfixExpression {} -> False
      LetExpression {} -> False
      LambdaExpression {} -> True
      IfExpression {} -> False
      CaseExpression {} -> False
      ParensExpression {} -> False
  
  --------------------------------------------------------------------------------
  -- Expression manipulators
  
  -- | Flatten an application f x y into (f,[x,y]).
  fargs :: Expression Type i l -> (Expression Type i l, [(Expression Type i l)])
  fargs e = go e []
    where
      go (ApplicationExpression _ f x) args = go f (x : args)
      go f args = (f, args)
  
  --------------------------------------------------------------------------------
  -- Substitutions
  
  substitute :: Eq i => i -> Expression Type i l -> Expression Type i l -> Expression Type i l
  substitute i arg = go
    where
      go =
        \case
          VariableExpression l i'
            | i == i' -> arg
            | otherwise -> VariableExpression l i'
          x@ConstructorExpression {} -> x
          x@ConstantExpression {} -> x
          ParensExpression _ e -> go e
          ApplicationExpression l f x -> ApplicationExpression l (go f) (go x)
          InfixExpression l x (s, f) y -> InfixExpression l (go x) (s, go f) (go y)
          LetExpression {} -> error "let expressions unsupported."
          CaseExpression l e cases ->
            CaseExpression l (go e) (map (\(CaseAlt l pat e') -> CaseAlt l pat (go e')) cases)
          IfExpression l a b c -> IfExpression l (go a) (go b) (go c)
          x@LiteralExpression {} -> x
          LambdaExpression l (Alternative l' args body) ->
            LambdaExpression l (Alternative l' args (go body))
  
  --------------------------------------------------------------------------------
  -- Lookups
  
  lookupName
    :: (MonadThrow m)
    => Name
    -> [BindGroup Type Name (Location)]
    -> m (Expression Type Name (Location))
  lookupName identifier binds =
    case listToMaybe (mapMaybe findIdent binds) of
      Nothing -> throwM (CouldntFindName identifier)
      Just i -> pure i
    where
      findIdent (BindGroup es is) =
        listToMaybe
          (mapMaybe
             (\case
                ImplicitlyTypedBinding _ (i, _) [Alternative _ [] e]
                  | i == identifier -> Just e
                _ -> Nothing)
             (concat is)) <|>
        listToMaybe
          (mapMaybe
             (\case
                ExplicitlyTypedBinding _ (i, _) _ [Alternative _ [] e]
                  | i == identifier -> Just e
                _ -> Nothing)
             es)
  
  lookupNameByString
    :: (MonadThrow m)
    => String
    -> [BindGroup Type Name (Location)]
    -> m (Expression Type Name (Location))
  lookupNameByString identifier binds =
    case listToMaybe (mapMaybe findIdent binds) of
      Nothing -> throwM (CouldntFindNameByString identifier)
      Just i -> pure i
    where
      findIdent (BindGroup es is) =
        listToMaybe
          (mapMaybe
             (\case
                ImplicitlyTypedBinding _ (ValueName _ i, _) [Alternative _ [] e]
                  | i == identifier -> Just e
                _ -> Nothing)
             (concat is)) <|>
        listToMaybe
          (mapMaybe
             (\case
                ExplicitlyTypedBinding _ (ValueName _ i, _) _ [Alternative _ [] e]
                  | i == identifier -> Just e
                _ -> Nothing)
             es)
  