Fix the implementation of the first rules

using State monad to keep track of total Nu count

Co-authored-by: Danila Danko <[email protected]>

eo-phi-normalizer/eo-phi-normalizer.cabal

@@ -54,6 +54,7 @@ library
     , base >=4.7 && <5
     , directory
     , filepath
+    , mtl
     , yaml
   default-language: Haskell2010
 
@@ -76,6 +77,7 @@ executable normalize-phi
     , directory
     , eo-phi-normalizer
     , filepath
+    , mtl
     , yaml
   default-language: Haskell2010
 
@@ -103,5 +105,6 @@ test-suite eo-phi-normalizer-test
     , filepath
     , hspec
     , hspec-discover
+    , mtl
     , yaml
   default-language: Haskell2010

eo-phi-normalizer/package.yaml

@@ -41,6 +41,7 @@ dependencies:
 - directory
 - filepath
 - yaml
+- mtl
 
 ghc-options:
 - -Wall

eo-phi-normalizer/src/Language/EO/Phi/Normalize.hs

@@ -1,20 +1,23 @@
 {-# LANGUAGE LambdaCase #-}
-{-# LANGUAGE RecordWildCards #-}
 
 module Language.EO.Phi.Normalize (
-  normalizeObjectWith,
+  normalizeObject,
   normalize,
   peelObject,
   unpeelObject,
 ) where
 
+import Control.Monad.State
 import Data.Maybe (fromMaybe)
-import Language.EO.Phi.Syntax.Abs
 import Numeric (showHex)
 
+import Control.Monad (forM)
+import Language.EO.Phi.Syntax.Abs
+
 data Context = Context
   { globalObject :: [Binding]
   , thisObject :: [Binding]
+  , totalNuCount :: Int
   }
 
 lookupBinding :: Attribute -> [Binding] -> Maybe Object
@@ -24,51 +27,60 @@ lookupBinding a (AlphaBinding a' object : bindings)
   | otherwise = lookupBinding a bindings
 lookupBinding _ _ = Nothing
 
+isNu :: Binding -> Bool
+isNu (AlphaBinding VTX _) = True
+isNu (EmptyBinding VTX) = True
+isNu _ = False
+
 -- | Normalize an input 𝜑-program.
 normalize :: Program -> Program
-normalize (Program bindings) = Program (map (normalizeBindingWith context) bindings)
+normalize (Program bindings) = evalState (Program <$> forM bindings normalizeBinding) context
  where
   context =
     Context
       { globalObject = bindings
       , thisObject = bindings
+      , totalNuCount = nuCount bindings
       }
+  nuCount binds = count isNu binds + sum (map (sum . map (nuCount . objectBindings) . values) binds)
+  count = (length .) . filter
+  values (AlphaBinding _ obj) = [obj]
+  values _ = []
+  objectBindings (Formation bs) = bs
+  objectBindings _ = []
 
-normalizeBindingWith :: Context -> Binding -> Binding
-normalizeBindingWith context = \case
-  AlphaBinding a object -> AlphaBinding a (normalizeObjectWith context object)
-  binding -> binding
+normalizeBinding :: Binding -> State Context Binding
+normalizeBinding = \case
+  AlphaBinding a object -> AlphaBinding a <$> normalizeObject object
+  binding -> pure binding
 
-count :: (a -> Bool) -> [a] -> Int
-count = (length .) . filter
+rule1 :: Object -> State Context Object
+rule1 (Formation bindings) = do
+  normalizedBindings <- forM bindings $ \case
+    AlphaBinding a object -> do
+      object' <- rule1 object
+      pure (AlphaBinding a object')
+    b -> pure b
+  finalBindings <-
+    if not $ any isNu normalizedBindings
+      then do
+        nus <- gets totalNuCount
+        modify (\c -> c{totalNuCount = totalNuCount c + 1})
+        let dataObject = Formation [DeltaBinding $ Bytes $ showHex nus ""]
+        pure (AlphaBinding VTX dataObject : normalizedBindings)
+      else do
+        pure normalizedBindings
+  pure (Formation finalBindings)
+rule1 object = pure object
 
-normalizeObjectWith :: Context -> Object -> Object
-normalizeObjectWith ctx@Context{..} object =
+normalizeObject :: Object -> State Context Object
+normalizeObject object = do
+  this <- gets thisObject
   case object of
     -- Rule 1
-    Formation bindings -> Formation bindings'
-     where
-      bindings'
-        | not $ any isNu bindings = AlphaBinding VTX (dataObject nu) : normalizedBindings
-        | otherwise = normalizedBindings
-      normalizedBindings = map (normalizeBindingWith ctx) bindings
-      nuCount binds = count isNu binds + sum (map (sum . map (nuCount . objectBindings) . values) binds)
-      dataObject n = Formation [DeltaBinding $ Bytes $ showHex n ""]
-
-      values (AlphaBinding _ obj) = [obj]
-      values _ = []
-
-      objectBindings (Formation bs) = bs
-      objectBindings _ = []
-
-      isNu (AlphaBinding VTX _) = True
-      isNu (EmptyBinding VTX) = True
-      isNu _ = False
-
-      nu = nuCount normalizedBindings
-    ThisDispatch a ->
-      fromMaybe object (lookupBinding a thisObject)
-    _ -> object
+    obj@(Formation _) -> rule1 obj
+    ThisDispatch a -> pure $ fromMaybe object (lookupBinding a this)
+    _ -> pure object
 
 -- | Split compound object into its head and applications/dispatch actions.
 peelObject :: Object -> PeeledObject