{- |
@.frame@: names for locals, replaced by their indices before the program is
parsed.
-}
module Wrench.Isa.Wasm32.LocalNames (
    resolveLocalNames,
) where

import Data.Char (isAlphaNum, isSpace)
import Data.List qualified as L
import Relude
import Relude.Extra (lookup)
import Wrench.Translator.Types (Preprocessor)

-- | In @.text@, from each label to the next one or the section's end, blank
-- the @.frame ($a, $b) $c@ line and write each name it lists as its index,
-- padded to the name's width. A name the @.frame@ does not list stays in
-- the source for the parser to reject.
resolveLocalNames :: Preprocessor
resolveLocalNames :: Preprocessor
resolveLocalNames String
src = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"\n" ([String] -> String)
-> Either (Int, Text) [String] -> Either (Int, Text) String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int
-> Bool
-> Map String Int
-> [String]
-> Either (Int, Text) [String]
forall {b}.
IsString b =>
Int
-> Bool -> Map String Int -> [String] -> Either (Int, b) [String]
go Int
0 Bool
False Map String Int
forall a. Monoid a => a
mempty (Char -> String -> [String]
splitOn Char
'\n' String
src)
    where
        go :: Int
-> Bool -> Map String Int -> [String] -> Either (Int, b) [String]
go Int
_ Bool
_ Map String Int
_ [] = [String] -> Either (Int, b) [String]
forall a b. b -> Either a b
Right []
        go Int
offset Bool
inText Map String Int
names (String
line : [String]
rest) = case String -> [String]
L.words (String -> String
code String
line) of
            String
".text" : [String]
_ -> (String
line String -> [String] -> [String]
forall a. a -> [a] -> [a]
:) ([String] -> [String])
-> Either (Int, b) [String] -> Either (Int, b) [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Map String Int -> Either (Int, b) [String]
next Bool
True Map String Int
forall a. Monoid a => a
mempty
            String
".data" : [String]
_ -> (String
line String -> [String] -> [String]
forall a. a -> [a] -> [a]
:) ([String] -> [String])
-> Either (Int, b) [String] -> Either (Int, b) [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Map String Int -> Either (Int, b) [String]
next Bool
False Map String Int
forall a. Monoid a => a
mempty
            String
w : [String]
_ | Bool
inText, String
":" String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`L.isSuffixOf` String
w -> (String
line String -> [String] -> [String]
forall a. a -> [a] -> [a]
:) ([String] -> [String])
-> Either (Int, b) [String] -> Either (Int, b) [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Map String Int -> Either (Int, b) [String]
next Bool
inText Map String Int
forall a. Monoid a => a
mempty
            [String]
_
                | Bool -> Bool
not Bool
inText -> (String
line String -> [String] -> [String]
forall a. a -> [a] -> [a]
:) ([String] -> [String])
-> Either (Int, b) [String] -> Either (Int, b) [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Map String Int -> Either (Int, b) [String]
next Bool
inText Map String Int
names
                | Just String
afterDirective <- String -> Maybe String
frameDirective String
line -> case String -> Maybe [String]
frameNames String
afterDirective of
                    Maybe [String]
Nothing ->
                        (Int, b) -> Either (Int, b) [String]
forall a b. a -> Either a b
Left (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
takeWhile Char -> Bool
isSpace String
line), b
"malformed .frame, expected distinct names as in .frame ($a, $b) $c")
                    Just [String]
declared -> ((Char -> Char) -> String -> String
forall a b. (a -> b) -> [a] -> [b]
map (Char -> Char -> Char
forall a b. a -> b -> a
const Char
' ') String
line String -> [String] -> [String]
forall a. a -> [a] -> [a]
:) ([String] -> [String])
-> Either (Int, b) [String] -> Either (Int, b) [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Map String Int -> Either (Int, b) [String]
next Bool
inText ([Item (Map String Int)] -> Map String Int
forall l. IsList l => [Item l] -> l
fromList ([String] -> [Int] -> [(String, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [String]
declared [Int
0 ..]))
                | Bool
otherwise -> (Map String Int -> String -> String
substitute Map String Int
names String
line String -> [String] -> [String]
forall a. a -> [a] -> [a]
:) ([String] -> [String])
-> Either (Int, b) [String] -> Either (Int, b) [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Map String Int -> Either (Int, b) [String]
next Bool
inText Map String Int
names
            where
                next :: Bool -> Map String Int -> Either (Int, b) [String]
next Bool
inText' Map String Int
names' = Int
-> Bool -> Map String Int -> [String] -> Either (Int, b) [String]
go (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
line Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Bool
inText' Map String Int
names' [String]
rest

-- | What follows @.frame@ on the line, if the line is one.
frameDirective :: String -> Maybe String
frameDirective :: String -> Maybe String
frameDirective String
line = case String -> String -> Maybe String
forall a. Eq a => [a] -> [a] -> Maybe [a]
L.stripPrefix String
".frame" ((Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace (String -> String
code String
line)) of
    Just s :: String
s@(Char
c : String
_) | Char -> Bool
isSpace Char
c Bool -> Bool -> Bool
|| Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'(' -> String -> Maybe String
forall a. a -> Maybe a
Just String
s
    Just [] -> String -> Maybe String
forall a. a -> Maybe a
Just []
    Maybe String
_ -> Maybe String
forall a. Maybe a
Nothing

-- | The names of @($a, $b) $c@: the parameters, then the reserved locals.
frameNames :: String -> Maybe [String]
frameNames :: String -> Maybe [String]
frameNames String
s = do
    '(' : inside <- String -> Maybe String
forall a. a -> Maybe a
Just ((Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace String
s)
    (params, _ : locals) <- Just (break (== ')') inside)
    declared <- (<>) <$> nameList params <*> nameList locals
    guard $ length (ordNub declared) == length declared
    return declared
    where
        nameList :: String -> Maybe [String]
nameList String
t
            | (Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Char -> Bool
isSpace String
t = [String] -> Maybe [String]
forall a. a -> Maybe a
Just []
            | Bool
otherwise = (String -> Maybe String) -> [String] -> Maybe [String]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM String -> Maybe String
name (Char -> String -> [String]
splitOn Char
',' String
t)
        name :: String -> Maybe String
name String
t = case (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
L.dropWhileEnd Char -> Bool
isSpace ((Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace String
t) of
            n :: String
n@(Char
'$' : Char
c : String
cs) | (Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Char -> Bool
isNameChar (Char
c Char -> String -> String
forall a. a -> [a] -> [a]
: String
cs) -> String -> Maybe String
forall a. a -> Maybe a
Just String
n
            String
_ -> Maybe String
forall a. Maybe a
Nothing

substitute :: Map String Int -> String -> String
substitute :: Map String Int -> String -> String
substitute Map String Int
names = \case
    Char
'$' : String
rest
        | (String
name, String
rest') <- (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
span Char -> Bool
isNameChar String
rest
        , Just Val (Map String Int)
idx <- Key (Map String Int)
-> Map String Int -> Maybe (Val (Map String Int))
forall t. StaticMap t => Key t -> t -> Maybe (Val t)
lookup (Char
'$' Char -> String -> String
forall a. a -> [a] -> [a]
: String
name) Map String Int
names ->
            let index :: String
index = Int -> String
forall b a. (Show a, IsString b) => a -> b
show Int
Val (Map String Int)
idx in String
index String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> Char -> String
forall a. Int -> a -> [a]
replicate (String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
name Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
index) Char
' ' String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Map String Int -> String -> String
substitute Map String Int
names String
rest'
    Char
c : String
rest -> Char
c Char -> String -> String
forall a. a -> [a] -> [a]
: Map String Int -> String -> String
substitute Map String Int
names String
rest
    [] -> []

isNameChar :: Char -> Bool
isNameChar :: Char -> Bool
isNameChar Char
c = Char -> Bool
isAlphaNum Char
c Bool -> Bool -> Bool
|| Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'_'

code :: String -> String
code :: String -> String
code = (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
';')

splitOn :: Char -> String -> [String]
splitOn :: Char -> String -> [String]
splitOn Char
sep String
s = case (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
break (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
sep) String
s of
    (String
chunk, []) -> [String
chunk]
    (String
chunk, Char
_ : String
rest) -> String
chunk String -> [String] -> [String]
forall a. a -> [a] -> [a]
: Char -> String -> [String]
splitOn Char
sep String
rest