{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE RecordWildCards #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

-- SPDX-FileCopyrightText: Copyright (c) 2025 Objectionary.com
-- SPDX-License-Identifier: MIT

module Rewriter (rewrite, rewrite', RewriteContext (..), defaultRewriteContext) where

import Ast
import Builder
import qualified Condition as C
import Control.Exception
import qualified Data.Map.Strict as M
import Data.Maybe (catMaybes, fromMaybe, isJust)
import Debug.Trace (trace)
import Logger (logDebug)
import Matcher (MetaValue (MvAttribute, MvBindings, MvExpression), Subst (Subst), combine, combineMany, defaultScope, matchProgram, substEmpty, substSingle)
import Misc (ensuredFile)
import Parser (parseProgram, parseProgramThrows)
import Pretty (PrintMode (SWEET), prettyAttribute, prettyExpression, prettyProgram, prettyProgram', prettySubsts)
import Replacer (replaceProgram)
import Text.Printf
import Yaml (ExtraArgument (..))
import qualified Yaml as Y

data RewriteContext = RewriteContext
  { RewriteContext -> Program
program :: Program,
    RewriteContext -> Integer
maxDepth :: Integer
  }

defaultRewriteContext :: Program -> RewriteContext
defaultRewriteContext :: Program -> RewriteContext
defaultRewriteContext Program
prog = Program -> Integer -> RewriteContext
RewriteContext Program
prog Integer
25

data RewriteException
  = CouldNotBuild {RewriteException -> Expression
expr :: Expression, RewriteException -> [Subst]
substs :: [Subst]}
  | CouldNotReplace {RewriteException -> Program
prog :: Program, RewriteException -> Expression
ptn :: Expression, RewriteException -> Expression
res :: Expression}
  deriving (Show RewriteException
Typeable RewriteException
(Typeable RewriteException, Show RewriteException) =>
(RewriteException -> SomeException)
-> (SomeException -> Maybe RewriteException)
-> (RewriteException -> String)
-> Exception RewriteException
SomeException -> Maybe RewriteException
RewriteException -> String
RewriteException -> SomeException
forall e.
(Typeable e, Show e) =>
(e -> SomeException)
-> (SomeException -> Maybe e) -> (e -> String) -> Exception e
$ctoException :: RewriteException -> SomeException
toException :: RewriteException -> SomeException
$cfromException :: SomeException -> Maybe RewriteException
fromException :: SomeException -> Maybe RewriteException
$cdisplayException :: RewriteException -> String
displayException :: RewriteException -> String
Exception)

instance Show RewriteException where
  show :: RewriteException -> String
show CouldNotBuild {[Subst]
Expression
expr :: RewriteException -> Expression
substs :: RewriteException -> [Subst]
expr :: Expression
substs :: [Subst]
..} =
    String -> String -> String -> String
forall r. PrintfType r => String -> r
printf
      String
"Couldn't build given expression with provided substitutions\n--Expression: %s\n--Substitutions: %s"
      (Expression -> String
prettyExpression Expression
expr)
      ([Subst] -> String
prettySubsts [Subst]
substs)
  show CouldNotReplace {Expression
Program
prog :: RewriteException -> Program
ptn :: RewriteException -> Expression
res :: RewriteException -> Expression
prog :: Program
ptn :: Expression
res :: Expression
..} =
    String -> String -> String -> String -> String
forall r. PrintfType r => String -> r
printf
      String
"Couldn't replace expression in program by pattern\nProgram: %s\n--Pattern: %s\n--Result: %s"
      (Program -> String
prettyProgram Program
prog)
      (Expression -> String
prettyExpression Expression
ptn)
      (Expression -> String
prettyExpression Expression
res)

-- Build pattern and result expression and replace patterns to results in given program
buildAndReplace :: Program -> Expression -> Expression -> [Subst] -> IO Program
buildAndReplace :: Program -> Expression -> Expression -> [Subst] -> IO Program
buildAndReplace Program
program Expression
ptn Expression
res [Subst]
substs =
  case (Expression -> [Subst] -> Maybe [(Expression, Expression)]
buildExpressions Expression
ptn [Subst]
substs, Expression -> [Subst] -> Maybe [(Expression, Expression)]
buildExpressions Expression
res [Subst]
substs) of
    (Just [(Expression, Expression)]
ptns, Just [(Expression, Expression)]
repls) ->
      let repls' :: [Expression]
repls' = ((Expression, Expression) -> Expression)
-> [(Expression, Expression)] -> [Expression]
forall a b. (a -> b) -> [a] -> [b]
map (Expression, Expression) -> Expression
forall a b. (a, b) -> a
fst [(Expression, Expression)]
repls
          ptns' :: [Expression]
ptns' = ((Expression, Expression) -> Expression)
-> [(Expression, Expression)] -> [Expression]
forall a b. (a -> b) -> [a] -> [b]
map (Expression, Expression) -> Expression
forall a b. (a, b) -> a
fst [(Expression, Expression)]
ptns
       in case Program -> [Expression] -> [Expression] -> Maybe Program
replaceProgram Program
program [Expression]
ptns' [Expression]
repls' of
            Just Program
prog -> Program -> IO Program
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Program
prog
            Maybe Program
_ -> RewriteException -> IO Program
forall e a. Exception e => e -> IO a
throwIO (Program -> Expression -> Expression -> RewriteException
CouldNotReplace Program
program Expression
ptn Expression
res)
    (Maybe [(Expression, Expression)]
Nothing, Maybe [(Expression, Expression)]
_) -> RewriteException -> IO Program
forall e a. Exception e => e -> IO a
throwIO (Expression -> [Subst] -> RewriteException
CouldNotBuild Expression
ptn [Subst]
substs)
    (Maybe [(Expression, Expression)]
_, Maybe [(Expression, Expression)]
Nothing) -> RewriteException -> IO Program
forall e a. Exception e => e -> IO a
throwIO (Expression -> [Subst] -> RewriteException
CouldNotBuild Expression
res [Subst]
substs)

-- Extend list of given substitutions with extra substitutions from 'where' yaml rule section
extraSubstitutions :: Program -> Maybe [Y.Extra] -> [Subst] -> [Subst]
extraSubstitutions :: Program -> Maybe [Extra] -> [Subst] -> [Subst]
extraSubstitutions Program
prog Maybe [Extra]
extras [Subst]
substs = case Maybe [Extra]
extras of
  Maybe [Extra]
Nothing -> [Subst]
substs
  Just [Extra]
extras' ->
    [Maybe Subst] -> [Subst]
forall a. [Maybe a] -> [a]
catMaybes
      [ (Maybe Subst -> Extra -> Maybe Subst)
-> Maybe Subst -> [Extra] -> Maybe Subst
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
          ( \(Just Subst
subst') Extra
extra -> do
              String
name <- case Extra -> ExtraArgument
Y.meta Extra
extra of
                ArgExpression (ExMeta String
name) -> String -> Maybe String
forall a. a -> Maybe a
Just String
name
                ArgAttribute (AtMeta String
name) -> String -> Maybe String
forall a. a -> Maybe a
Just String
name
                ArgBinding (BiMeta String
name) -> String -> Maybe String
forall a. a -> Maybe a
Just String
name
                ExtraArgument
_ -> Maybe String
forall a. Maybe a
Nothing
              let func :: String
func = Extra -> String
Y.function Extra
extra
                  args :: [ExtraArgument]
args = Extra -> [ExtraArgument]
Y.args Extra
extra
              Term
term <- String -> [ExtraArgument] -> Subst -> Program -> Maybe Term
buildTermFromFunction String
func [ExtraArgument]
args Subst
subst' Program
prog
              let meta :: MetaValue
meta = case Term
term of
                    TeExpression Expression
expr -> Expression -> Expression -> MetaValue
MvExpression Expression
expr Expression
defaultScope
                    TeAttribute Attribute
attr -> Attribute -> MetaValue
MvAttribute Attribute
attr
              Subst -> Subst -> Maybe Subst
combine (String -> MetaValue -> Subst
substSingle String
name MetaValue
meta) Subst
subst'
          )
          (Subst -> Maybe Subst
forall a. a -> Maybe a
Just Subst
subst)
          [Extra]
extras'
        | Subst
subst <- [Subst]
substs
      ]

rewrite :: Program -> Program -> [Y.Rule] -> IO Program
rewrite :: Program -> Program -> [Rule] -> IO Program
rewrite Program
program Program
_ [] = Program -> IO Program
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Program
program
rewrite Program
program Program
program' (Rule
rule : [Rule]
rest) = do
  let ptn :: Expression
ptn = Rule -> Expression
Y.pattern Rule
rule
      res :: Expression
res = Rule -> Expression
Y.result Rule
rule
      condition :: Maybe Condition
condition = Rule -> Maybe Condition
Y.when Rule
rule
  Program
prog <- case Expression -> Maybe Condition -> Program -> Maybe [Subst]
C.matchProgramWithCondition Expression
ptn Maybe Condition
condition Program
program of
    Maybe [Subst]
Nothing -> Program -> IO Program
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Program
program
    Just [Subst]
matched -> do
      let substs :: [Subst]
substs = Program -> Maybe [Extra] -> [Subst] -> [Subst]
extraSubstitutions Program
program' (Rule -> Maybe [Extra]
Y.where_ Rule
rule) [Subst]
matched
      Program
prog' <- Program -> Expression -> Expression -> [Subst] -> IO Program
buildAndReplace Program
program Expression
ptn Expression
res [Subst]
substs
      String -> IO ()
logDebug (String -> String -> String -> String
forall r. PrintfType r => String -> r
printf String
"%s\n%s" (String -> Maybe String -> String
forall a. a -> Maybe a -> a
fromMaybe String
"unknown" (Rule -> Maybe String
Y.name Rule
rule)) (Program -> PrintMode -> String
prettyProgram' Program
prog' PrintMode
SWEET))
      Program -> IO Program
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Program
prog'
  Program -> Program -> [Rule] -> IO Program
rewrite Program
prog Program
program' [Rule]
rest

-- @todo #169:30min Memorize previous rewritten programs. Right now in order not to
--  get an infinite recursion during rewriting we just count have many times we apply
--  rewriting rules. If we reach given amount - we just stop. It's not idiomatic and may
--  not work on big programs. We need to introduce some mechanism which would memorize
--  all rewritten program on each step and if on some step we get the program that have already
--  been memorized - we fail because we got into infinite recursion. Ofc we should keep counting
--  rewriting cycles if program just only grows on each rewriting.
rewrite' :: Program -> [Y.Rule] -> RewriteContext -> IO Program
rewrite' :: Program -> [Rule] -> RewriteContext -> IO Program
rewrite' Program
prog [Rule]
rules RewriteContext {Integer
Program
program :: RewriteContext -> Program
maxDepth :: RewriteContext -> Integer
program :: Program
maxDepth :: Integer
..} = Program -> Integer -> IO Program
_rewrite Program
prog Integer
0
  where
    _rewrite :: Program -> Integer -> IO Program
    _rewrite :: Program -> Integer -> IO Program
_rewrite Program
prog Integer
count = do
      String -> IO ()
logDebug (String -> Integer -> Integer -> String
forall r. PrintfType r => String -> r
printf String
"Starting rewriting cycle %d out of %d" Integer
count Integer
maxDepth)
      if Integer
count Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
maxDepth
        then do
          String -> IO ()
logDebug (String -> String
forall r. PrintfType r => String -> r
printf String
"Max amount of rewriting cycles is reached, rewriting is stopped")
          Program -> IO Program
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Program
prog
        else do
          Program
rewritten <- Program -> Program -> [Rule] -> IO Program
rewrite Program
prog Program
program [Rule]
rules
          if Program
rewritten Program -> Program -> Bool
forall a. Eq a => a -> a -> Bool
== Program
prog
            then do
              String -> IO ()
logDebug String
"Rewriting is stopped since it does not affect program anymore"
              Program -> IO Program
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Program
rewritten
            else Program -> Integer -> IO Program
_rewrite Program
rewritten (Integer
count Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)