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

module Replacer
  ( replaceExpression
  , replaceExpressionFast
  , ReplaceExpressionFunc
  )
where

import AST
import Data.List (isInfixOf, isPrefixOf)

type ReplaceState a = (a, [Expression], [Expression -> Expression])

type ReplaceExpressionFunc' = ReplaceState Expression -> ReplaceState Expression

type ReplaceExpressionFunc = ReplaceState Expression -> Expression

replaceBindings :: ReplaceState [Binding] -> ReplaceExpressionFunc' -> ReplaceState [Binding]
replaceBindings :: ReplaceState [Binding]
-> ReplaceExpressionFunc' -> ReplaceState [Binding]
replaceBindings state :: ReplaceState [Binding]
state@([Binding]
_, [], [Expression -> Expression]
_) ReplaceExpressionFunc'
_ = ReplaceState [Binding]
state
replaceBindings state :: ReplaceState [Binding]
state@([Binding]
_, [Expression]
_, []) ReplaceExpressionFunc'
_ = ReplaceState [Binding]
state
replaceBindings state :: ReplaceState [Binding]
state@([], [Expression]
_, [Expression -> Expression]
_) ReplaceExpressionFunc'
_ = ReplaceState [Binding]
state
replaceBindings (BiTau Attribute
attr Expression
expr : [Binding]
bds, [Expression]
ptns, [Expression -> Expression]
repls) ReplaceExpressionFunc'
func =
  let (Expression
expr', [Expression]
ptns', [Expression -> Expression]
repls') = ReplaceExpressionFunc'
func (Expression
expr, [Expression]
ptns, [Expression -> Expression]
repls)
      ([Binding]
bds', [Expression]
ptns'', [Expression -> Expression]
repls'') = ReplaceState [Binding]
-> ReplaceExpressionFunc' -> ReplaceState [Binding]
replaceBindings ([Binding]
bds, [Expression]
ptns', [Expression -> Expression]
repls') ReplaceExpressionFunc'
func
   in (Attribute -> Expression -> Binding
BiTau Attribute
attr Expression
expr' Binding -> [Binding] -> [Binding]
forall a. a -> [a] -> [a]
: [Binding]
bds', [Expression]
ptns'', [Expression -> Expression]
repls'')
replaceBindings (Binding
bd : [Binding]
bds, [Expression]
ptns, [Expression -> Expression]
repls) ReplaceExpressionFunc'
func =
  let ([Binding]
bds', [Expression]
ptns', [Expression -> Expression]
repls') = ReplaceState [Binding]
-> ReplaceExpressionFunc' -> ReplaceState [Binding]
replaceBindings ([Binding]
bds, [Expression]
ptns, [Expression -> Expression]
repls) ReplaceExpressionFunc'
func
   in (Binding
bd Binding -> [Binding] -> [Binding]
forall a. a -> [a] -> [a]
: [Binding]
bds', [Expression]
ptns', [Expression -> Expression]
repls')

replaceArgument :: ReplaceState Argument -> ReplaceExpressionFunc' -> ReplaceState Argument
replaceArgument :: ReplaceState Argument
-> ReplaceExpressionFunc' -> ReplaceState Argument
replaceArgument (ArTau Attribute
attr Expression
expr, [Expression]
ptns, [Expression -> Expression]
repls) ReplaceExpressionFunc'
func =
  let (Expression
expr', [Expression]
ptns', [Expression -> Expression]
repls') = ReplaceExpressionFunc'
func (Expression
expr, [Expression]
ptns, [Expression -> Expression]
repls)
   in (Attribute -> Expression -> Argument
ArTau Attribute
attr Expression
expr', [Expression]
ptns', [Expression -> Expression]
repls')
replaceArgument (ArAlpha Alpha
alpha Expression
expr, [Expression]
ptns, [Expression -> Expression]
repls) ReplaceExpressionFunc'
func =
  let (Expression
expr', [Expression]
ptns', [Expression -> Expression]
repls') = ReplaceExpressionFunc'
func (Expression
expr, [Expression]
ptns, [Expression -> Expression]
repls)
   in (Alpha -> Expression -> Argument
ArAlpha Alpha
alpha Expression
expr', [Expression]
ptns', [Expression -> Expression]
repls')

replaceExpression' :: ReplaceExpressionFunc'
replaceExpression' :: ReplaceExpressionFunc'
replaceExpression' state :: (Expression, [Expression], [Expression -> Expression])
state@(Expression
expr, ptns :: [Expression]
ptns@(Expression
ptn : [Expression]
_ptns), repls :: [Expression -> Expression]
repls@(Expression -> Expression
repl : [Expression -> Expression]
_repls))
  | Expression -> Bool
inert Expression
expr Bool -> Bool -> Bool
&& Bool -> Bool
not (Expression -> Bool
inert Expression
ptn) = (Expression, [Expression], [Expression -> Expression])
state
  | Expression
expr Expression -> Expression -> Bool
forall a. Eq a => a -> a -> Bool
== Expression
ptn = let ([Expression]
ptns', [Expression -> Expression]
repls') = Expression
-> ([Expression], [Expression -> Expression])
-> ([Expression], [Expression -> Expression])
swallowed Expression
expr ([Expression]
_ptns, [Expression -> Expression]
_repls) in ReplaceExpressionFunc'
replaceExpression' (Expression -> Expression
repl Expression
expr, [Expression]
ptns', [Expression -> Expression]
repls')
  | Bool
otherwise = case Expression
expr of
      ExDispatch Expression
inner Attribute
attr ->
        let (Expression
expr', [Expression]
ptns', [Expression -> Expression]
repls') = ReplaceExpressionFunc'
replaceExpression' (Expression
inner, [Expression]
ptns, [Expression -> Expression]
repls)
         in (Expression -> Attribute -> Expression
ExDispatch Expression
expr' Attribute
attr, [Expression]
ptns', [Expression -> Expression]
repls')
      ExApplication Expression
inner Argument
arg ->
        let (Expression
expr', [Expression]
ptns', [Expression -> Expression]
repls') = ReplaceExpressionFunc'
replaceExpression' (Expression
inner, [Expression]
ptns, [Expression -> Expression]
repls)
            (Argument
arg', [Expression]
ptns'', [Expression -> Expression]
repls'') = ReplaceState Argument
-> ReplaceExpressionFunc' -> ReplaceState Argument
replaceArgument (Argument
arg, [Expression]
ptns', [Expression -> Expression]
repls') ReplaceExpressionFunc'
replaceExpression'
         in (Expression -> Argument -> Expression
ExApplication Expression
expr' Argument
arg', [Expression]
ptns'', [Expression -> Expression]
repls'')
      ExFormation [Binding]
bds ->
        let ([Binding]
bds', [Expression]
ptns', [Expression -> Expression]
repls') = ReplaceState [Binding]
-> ReplaceExpressionFunc' -> ReplaceState [Binding]
replaceBindings ([Binding]
bds, [Expression]
ptns, [Expression -> Expression]
repls) ReplaceExpressionFunc'
replaceExpression'
         in ([Binding] -> Expression
ExFormation [Binding]
bds', [Expression]
ptns', [Expression -> Expression]
repls')
      Expression
_ -> (Expression, [Expression], [Expression -> Expression])
state
replaceExpression' (Expression, [Expression], [Expression -> Expression])
state = (Expression, [Expression], [Expression -> Expression])
state

replaceBindingsFast :: Expression -> ReplaceState [Binding] -> ReplaceState [Binding]
replaceBindingsFast :: Expression -> ReplaceState [Binding] -> ReplaceState [Binding]
replaceBindingsFast Expression
_ state :: ReplaceState [Binding]
state@([Binding]
_, [], [Expression -> Expression]
_) = ReplaceState [Binding]
state
replaceBindingsFast Expression
_ state :: ReplaceState [Binding]
state@([Binding]
_, [Expression]
_, []) = ReplaceState [Binding]
state
replaceBindingsFast Expression
expr ([Binding]
bds, Expression
ptn : [Expression]
ptns, Expression -> Expression
repl : [Expression -> Expression]
repls) = case (Expression
ptn, Expression -> Expression
repl Expression
expr) of
  (ExFormation [], ExFormation [Binding]
rbds) -> Expression -> ReplaceState [Binding] -> ReplaceState [Binding]
replaceBindingsFast Expression
expr ([Binding]
rbds, [Expression]
ptns, [Expression -> Expression]
repls)
  (ExFormation [Binding]
pbds, ExFormation [Binding]
rbds)
    | [Binding]
pbds [Binding] -> [Binding] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isInfixOf` [Binding]
bds -> Expression -> ReplaceState [Binding] -> ReplaceState [Binding]
replaceBindingsFast Expression
expr ([Binding] -> [Binding] -> [Binding] -> [Binding]
findAndReplace [Binding]
bds [Binding]
pbds [Binding]
rbds, [Expression]
ptns, [Expression -> Expression]
repls)
  (Expression, Expression)
_ ->
    let ([Binding]
bds', [Expression]
ptns', [Expression -> Expression]
repls') = Expression -> ReplaceState [Binding] -> ReplaceState [Binding]
replaceBindingsFast Expression
expr ([Binding]
bds, [Expression]
ptns, [Expression -> Expression]
repls)
     in ([Binding]
bds', Expression
ptn Expression -> [Expression] -> [Expression]
forall a. a -> [a] -> [a]
: [Expression]
ptns', Expression -> Expression
repl (Expression -> Expression)
-> [Expression -> Expression] -> [Expression -> Expression]
forall a. a -> [a] -> [a]
: [Expression -> Expression]
repls')
  where
    findAndReplace :: [Binding] -> [Binding] -> [Binding] -> [Binding]
    findAndReplace :: [Binding] -> [Binding] -> [Binding] -> [Binding]
findAndReplace [] [Binding]
_ [Binding]
_ = []
    findAndReplace [Binding]
_ [] [Binding]
rbds = [Binding]
rbds
    findAndReplace xs :: [Binding]
xs@(Binding
x : [Binding]
xs') [Binding]
pbds [Binding]
rbds
      | [Binding]
pbds [Binding] -> [Binding] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` [Binding]
xs = [Binding]
rbds [Binding] -> [Binding] -> [Binding]
forall a. [a] -> [a] -> [a]
++ [Binding] -> [Binding] -> [Binding] -> [Binding]
findAndReplace (Int -> [Binding] -> [Binding]
forall a. Int -> [a] -> [a]
drop ([Binding] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Binding]
pbds) [Binding]
xs) [Binding]
pbds [Binding]
rbds
      | Bool
otherwise = Binding
x Binding -> [Binding] -> [Binding]
forall a. a -> [a] -> [a]
: [Binding] -> [Binding] -> [Binding] -> [Binding]
findAndReplace [Binding]
xs' [Binding]
pbds [Binding]
rbds

replaceExpressionFast' :: ReplaceExpressionFunc'
replaceExpressionFast' :: ReplaceExpressionFunc'
replaceExpressionFast' state :: (Expression, [Expression], [Expression -> Expression])
state@(Expression
_, [], [Expression -> Expression]
_) = (Expression, [Expression], [Expression -> Expression])
state
replaceExpressionFast' state :: (Expression, [Expression], [Expression -> Expression])
state@(Expression
_, [Expression]
_, []) = (Expression, [Expression], [Expression -> Expression])
state
replaceExpressionFast' state :: (Expression, [Expression], [Expression -> Expression])
state@(Expression
expr, [Expression]
ptns, [Expression -> Expression]
repls) = case Expression
expr of
  ExFormation [Binding]
bds ->
    let ([Binding]
bds', [Expression]
ptns', [Expression -> Expression]
repls') = ReplaceState [Binding]
-> ReplaceExpressionFunc' -> ReplaceState [Binding]
replaceBindings (Expression -> ReplaceState [Binding] -> ReplaceState [Binding]
replaceBindingsFast Expression
expr ([Binding]
bds, [Expression]
ptns, [Expression -> Expression]
repls)) ReplaceExpressionFunc'
replaceExpressionFast'
     in ([Binding] -> Expression
ExFormation [Binding]
bds', [Expression]
ptns', [Expression -> Expression]
repls')
  ExDispatch Expression
inner Attribute
attr ->
    let (Expression
expr', [Expression]
ptns', [Expression -> Expression]
repls') = ReplaceExpressionFunc'
replaceExpressionFast' (Expression
inner, [Expression]
ptns, [Expression -> Expression]
repls)
     in (Expression -> Attribute -> Expression
ExDispatch Expression
expr' Attribute
attr, [Expression]
ptns', [Expression -> Expression]
repls')
  ExApplication Expression
inner Argument
arg ->
    let (Expression
expr', [Expression]
ptns', [Expression -> Expression]
repls') = ReplaceExpressionFunc'
replaceExpressionFast' (Expression
inner, [Expression]
ptns, [Expression -> Expression]
repls)
        (Argument
arg', [Expression]
ptns'', [Expression -> Expression]
repls'') = ReplaceState Argument
-> ReplaceExpressionFunc' -> ReplaceState Argument
replaceArgument (Argument
arg, [Expression]
ptns', [Expression -> Expression]
repls') ReplaceExpressionFunc'
replaceExpressionFast'
     in (Expression -> Argument -> Expression
ExApplication Expression
expr' Argument
arg', [Expression]
ptns'', [Expression -> Expression]
repls'')
  Expression
_ -> (Expression, [Expression], [Expression -> Expression])
state

swallowed :: Expression -> ([Expression], [Expression -> Expression]) -> ([Expression], [Expression -> Expression])
swallowed :: Expression
-> ([Expression], [Expression -> Expression])
-> ([Expression], [Expression -> Expression])
swallowed Expression
replaced = [Expression]
-> ([Expression], [Expression -> Expression])
-> ([Expression], [Expression -> Expression])
go []
  where
    go :: [Expression] -> ([Expression], [Expression -> Expression]) -> ([Expression], [Expression -> Expression])
    go :: [Expression]
-> ([Expression], [Expression -> Expression])
-> ([Expression], [Expression -> Expression])
go [Expression]
dropped (Expression
ptn : [Expression]
ptns, Expression -> Expression
_ : [Expression -> Expression]
repls)
      | [Expression] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((Expression -> Bool) -> [Expression] -> [Expression]
forall a. (a -> Bool) -> [a] -> [a]
filter (Expression -> Expression -> Bool
forall a. Eq a => a -> a -> Bool
== Expression
ptn) [Expression]
dropped) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Expression -> Int) -> [Expression] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Expression -> Expression -> Int
occurrences Expression
ptn) (Expression -> [Expression]
children Expression
replaced)) = [Expression]
-> ([Expression], [Expression -> Expression])
-> ([Expression], [Expression -> Expression])
go (Expression
ptn Expression -> [Expression] -> [Expression]
forall a. a -> [a] -> [a]
: [Expression]
dropped) ([Expression]
ptns, [Expression -> Expression]
repls)
    go [Expression]
_ ([Expression], [Expression -> Expression])
pending = ([Expression], [Expression -> Expression])
pending
    occurrences :: Expression -> Expression -> Int
    occurrences :: Expression -> Expression -> Int
occurrences Expression
ptn Expression
term = (if Expression
term Expression -> Expression -> Bool
forall a. Eq a => a -> a -> Bool
== Expression
ptn then Int
1 else Int
0) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Expression -> Int) -> [Expression] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Expression -> Expression -> Int
occurrences Expression
ptn) (Expression -> [Expression]
children Expression
term))
    children :: Expression -> [Expression]
    children :: Expression -> [Expression]
children (ExFormation [Binding]
bds) = [Expression
inner | BiTau Attribute
_ Expression
inner <- [Binding]
bds]
    children (ExDispatch Expression
inner Attribute
_) = [Expression
inner]
    children (ExApplication Expression
inner (ArTau Attribute
_ Expression
arg)) = [Expression
inner, Expression
arg]
    children (ExApplication Expression
inner (ArAlpha Alpha
_ Expression
arg)) = [Expression
inner, Expression
arg]
    children (ExPhiMeet Maybe String
_ Int
_ Expression
inner) = [Expression
inner]
    children (ExPhiAgain Maybe String
_ Int
_ Expression
inner) = [Expression
inner]
    children Expression
_ = []

replaceExpression :: ReplaceExpressionFunc
replaceExpression :: ReplaceExpressionFunc
replaceExpression (Expression, [Expression], [Expression -> Expression])
state =
  let (Expression
expr, [Expression]
_, [Expression -> Expression]
_) = ReplaceExpressionFunc'
replaceExpression' (Expression, [Expression], [Expression -> Expression])
state
   in Expression
expr

replaceExpressionFast :: ReplaceExpressionFunc
replaceExpressionFast :: ReplaceExpressionFunc
replaceExpressionFast (Expression, [Expression], [Expression -> Expression])
state =
  let (Expression
expr, [Expression]
_, [Expression -> Expression]
_) = ReplaceExpressionFunc'
replaceExpressionFast' (Expression, [Expression], [Expression -> Expression])
state
   in Expression
expr