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

module Canonizer (canonize, canonizeExpr, lambdaNames) where

import AST
import qualified Data.Map as Map
import qualified Data.Text as T
import Rewriter (Rewritten)

type CanonState = (Map.Map T.Text T.Text, Int)

canonizeBindings :: [Binding] -> CanonState -> ([Binding], CanonState)
canonizeBindings :: [Binding] -> CanonState -> ([Binding], CanonState)
canonizeBindings [] CanonState
st = ([], CanonState
st)
canonizeBindings ((BiLambda (Function Text
name)) : [Binding]
rest) CanonState
st
  | Text
name Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Text
T.pack String
"Package" =
      let ([Binding]
bds', CanonState
st') = [Binding] -> CanonState -> ([Binding], CanonState)
canonizeBindings [Binding]
rest CanonState
st
       in (Function -> Binding
BiLambda (Text -> Function
Function Text
name) Binding -> [Binding] -> [Binding]
forall a. a -> [a] -> [a]
: [Binding]
bds', CanonState
st')
canonizeBindings ((BiLambda (Function Text
name)) : [Binding]
rest) (Map Text Text
names, Int
idx) =
  case Text -> Map Text Text -> Maybe Text
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
name Map Text Text
names of
    Just Text
canonical ->
      let ([Binding]
bds', CanonState
st') = [Binding] -> CanonState -> ([Binding], CanonState)
canonizeBindings [Binding]
rest (Map Text Text
names, Int
idx)
       in (Function -> Binding
BiLambda (Text -> Function
Function Text
canonical) Binding -> [Binding] -> [Binding]
forall a. a -> [a] -> [a]
: [Binding]
bds', CanonState
st')
    Maybe Text
Nothing ->
      let canonical :: Text
canonical = String -> Text
T.pack (String
"Fn" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
idx)
          ([Binding]
bds', CanonState
st') = [Binding] -> CanonState -> ([Binding], CanonState)
canonizeBindings [Binding]
rest (Text -> Text -> Map Text Text -> Map Text Text
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Text
name Text
canonical Map Text Text
names, Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
       in (Function -> Binding
BiLambda (Text -> Function
Function Text
canonical) Binding -> [Binding] -> [Binding]
forall a. a -> [a] -> [a]
: [Binding]
bds', CanonState
st')
canonizeBindings (BiTau Attribute
attr Expression
expr : [Binding]
rest) CanonState
st =
  let (Expression
expr', CanonState
st') = Expression -> CanonState -> (Expression, CanonState)
canonizeExpression Expression
expr CanonState
st
      ([Binding]
bds', CanonState
st'') = [Binding] -> CanonState -> ([Binding], CanonState)
canonizeBindings [Binding]
rest CanonState
st'
   in (Attribute -> Expression -> Binding
BiTau Attribute
attr Expression
expr' Binding -> [Binding] -> [Binding]
forall a. a -> [a] -> [a]
: [Binding]
bds', CanonState
st'')
canonizeBindings (Binding
bd : [Binding]
rest) CanonState
st =
  let ([Binding]
bds', CanonState
st') = [Binding] -> CanonState -> ([Binding], CanonState)
canonizeBindings [Binding]
rest CanonState
st
   in (Binding
bd Binding -> [Binding] -> [Binding]
forall a. a -> [a] -> [a]
: [Binding]
bds', CanonState
st')

canonizeExpression :: Expression -> CanonState -> (Expression, CanonState)
canonizeExpression :: Expression -> CanonState -> (Expression, CanonState)
canonizeExpression (ExFormation [Binding]
bds) CanonState
st =
  let ([Binding]
bds', CanonState
st') = [Binding] -> CanonState -> ([Binding], CanonState)
canonizeBindings [Binding]
bds CanonState
st
   in ([Binding] -> Expression
ExFormation [Binding]
bds', CanonState
st')
canonizeExpression (ExDispatch Expression
expr Attribute
attr) CanonState
st =
  let (Expression
expr', CanonState
st') = Expression -> CanonState -> (Expression, CanonState)
canonizeExpression Expression
expr CanonState
st
   in (Expression -> Attribute -> Expression
ExDispatch Expression
expr' Attribute
attr, CanonState
st')
canonizeExpression (ExApplication Expression
expr Argument
arg) CanonState
st =
  let (Expression
expr', CanonState
st') = Expression -> CanonState -> (Expression, CanonState)
canonizeExpression Expression
expr CanonState
st
      (Argument
arg', CanonState
st'') = Argument -> CanonState -> (Argument, CanonState)
canonizeArgument Argument
arg CanonState
st'
   in (Expression -> Argument -> Expression
ExApplication Expression
expr' Argument
arg', CanonState
st'')
canonizeExpression (ExPhiMeet Maybe String
prefix Int
num Expression
expr) CanonState
st =
  let (Expression
expr', CanonState
st') = Expression -> CanonState -> (Expression, CanonState)
canonizeExpression Expression
expr CanonState
st
   in (Maybe String -> Int -> Expression -> Expression
ExPhiMeet Maybe String
prefix Int
num Expression
expr', CanonState
st')
canonizeExpression (ExPhiAgain Maybe String
prefix Int
num Expression
expr) CanonState
st =
  let (Expression
expr', CanonState
st') = Expression -> CanonState -> (Expression, CanonState)
canonizeExpression Expression
expr CanonState
st
   in (Maybe String -> Int -> Expression -> Expression
ExPhiAgain Maybe String
prefix Int
num Expression
expr', CanonState
st')
canonizeExpression Expression
expr CanonState
st = (Expression
expr, CanonState
st)

canonizeArgument :: Argument -> CanonState -> (Argument, CanonState)
canonizeArgument :: Argument -> CanonState -> (Argument, CanonState)
canonizeArgument (ArTau Attribute
attr Expression
expr) CanonState
st =
  let (Expression
expr', CanonState
st') = Expression -> CanonState -> (Expression, CanonState)
canonizeExpression Expression
expr CanonState
st
   in (Attribute -> Expression -> Argument
ArTau Attribute
attr Expression
expr', CanonState
st')
canonizeArgument (ArAlpha Alpha
alpha Expression
expr) CanonState
st =
  let (Expression
expr', CanonState
st') = Expression -> CanonState -> (Expression, CanonState)
canonizeExpression Expression
expr CanonState
st
   in (Alpha -> Expression -> Argument
ArAlpha Alpha
alpha Expression
expr', CanonState
st')

lambdaNames :: Expression -> [T.Text]
lambdaNames :: Expression -> [Text]
lambdaNames (ExFormation [Binding]
bds) = (Binding -> [Text]) -> [Binding] -> [Text]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Binding -> [Text]
named [Binding]
bds
  where
    named :: Binding -> [T.Text]
    named :: Binding -> [Text]
named (BiLambda (Function Text
name)) = [Text
name]
    named (BiTau Attribute
_ Expression
expr) = Expression -> [Text]
lambdaNames Expression
expr
    named Binding
_ = []
lambdaNames (ExDispatch Expression
expr Attribute
_) = Expression -> [Text]
lambdaNames Expression
expr
lambdaNames (ExApplication Expression
expr (ArTau Attribute
_ Expression
arg)) = Expression -> [Text]
lambdaNames Expression
expr [Text] -> [Text] -> [Text]
forall a. [a] -> [a] -> [a]
++ Expression -> [Text]
lambdaNames Expression
arg
lambdaNames (ExApplication Expression
expr (ArAlpha Alpha
_ Expression
arg)) = Expression -> [Text]
lambdaNames Expression
expr [Text] -> [Text] -> [Text]
forall a. [a] -> [a] -> [a]
++ Expression -> [Text]
lambdaNames Expression
arg
lambdaNames (ExPhiMeet Maybe String
_ Int
_ Expression
expr) = Expression -> [Text]
lambdaNames Expression
expr
lambdaNames (ExPhiAgain Maybe String
_ Int
_ Expression
expr) = Expression -> [Text]
lambdaNames Expression
expr
lambdaNames Expression
_ = []

canonizeExpr :: Expression -> Expression
canonizeExpr :: Expression -> Expression
canonizeExpr Expression
expr = (Expression, CanonState) -> Expression
forall a b. (a, b) -> a
fst (Expression -> CanonState -> (Expression, CanonState)
canonizeExpression Expression
expr (Map Text Text
forall k a. Map k a
Map.empty, Int
1))

canonize :: [Rewritten] -> [Rewritten]
canonize :: [Rewritten] -> [Rewritten]
canonize [] = []
canonize ((Expression
expr, Maybe (Judgment, String)
maybeRule) : [Rewritten]
rest) = (Expression -> Expression
canonizeExpr Expression
expr, Maybe (Judgment, String)
maybeRule) Rewritten -> [Rewritten] -> [Rewritten]
forall a. a -> [a] -> [a]
: [Rewritten] -> [Rewritten]
canonize [Rewritten]
rest