-- \* Basic utilities for interpreting tuples as records.
module Language.Futhark.Tuple
  ( areTupleFields,
    tupleFields,
    tupleFieldNames,
    sortFields,
  )
where

import Data.Char (isDigit, ord)
import Data.List (sortOn)
import Data.Map qualified as M
import Data.Text qualified as T
import Language.Futhark.Core (Name, nameFromString, nameToText)

-- | Does this record map correspond to a tuple?
areTupleFields :: M.Map Name a -> Maybe [a]
areTupleFields :: forall a. Map Name a -> Maybe [a]
areTupleFields Map Name a
fs =
  let fs' :: [(Name, a)]
fs' = Map Name a -> [(Name, a)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name a
fs
   in if (Map Name a -> Bool
forall a. Map Name a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Map Name a
fs Bool -> Bool -> Bool
|| [(Name, a)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Name, a)]
fs' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1)
        Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Name -> Name -> Bool) -> [Name] -> [Name] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
(==) (((Name, a) -> Name) -> [(Name, a)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, a) -> Name
forall a b. (a, b) -> a
fst [(Name, a)]
fs') [Name]
tupleFieldNames)
        then [a] -> Maybe [a]
forall a. a -> Maybe a
Just ([a] -> Maybe [a]) -> [a] -> Maybe [a]
forall a b. (a -> b) -> a -> b
$ ((Name, a) -> a) -> [(Name, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Name, a) -> a
forall a b. (a, b) -> b
snd [(Name, a)]
fs'
        else Maybe [a]
forall a. Maybe a
Nothing

-- | Construct a record map corresponding to a tuple.
tupleFields :: [a] -> M.Map Name a
tupleFields :: forall a. [a] -> Map Name a
tupleFields [a]
as = [(Name, a)] -> Map Name a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, a)] -> Map Name a) -> [(Name, a)] -> Map Name a
forall a b. (a -> b) -> a -> b
$ [Name] -> [a] -> [(Name, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tupleFieldNames [a]
as

-- | Increasing field names for a tuple (starts at 0).
tupleFieldNames :: [Name]
tupleFieldNames :: [Name]
tupleFieldNames = (Int -> Name) -> [Int] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (String -> Name
nameFromString (String -> Name) -> (Int -> String) -> Int -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> String
forall a. Show a => a -> String
show) [(Int
0 :: Int) ..]

-- | Sort fields by their name; taking care to sort numeric fields by
-- their numeric value.  This ensures that tuples and tuple-like
-- records match.
sortFields :: M.Map Name a -> [(Name, a)]
sortFields :: forall a. Map Name a -> [(Name, a)]
sortFields Map Name a
l = ((Either Int Name, (Name, a)) -> (Name, a))
-> [(Either Int Name, (Name, a))] -> [(Name, a)]
forall a b. (a -> b) -> [a] -> [b]
map (Either Int Name, (Name, a)) -> (Name, a)
forall a b. (a, b) -> b
snd ([(Either Int Name, (Name, a))] -> [(Name, a)])
-> [(Either Int Name, (Name, a))] -> [(Name, a)]
forall a b. (a -> b) -> a -> b
$ ((Either Int Name, (Name, a)) -> Either Int Name)
-> [(Either Int Name, (Name, a))] -> [(Either Int Name, (Name, a))]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Either Int Name, (Name, a)) -> Either Int Name
forall a b. (a, b) -> a
fst ([(Either Int Name, (Name, a))] -> [(Either Int Name, (Name, a))])
-> [(Either Int Name, (Name, a))] -> [(Either Int Name, (Name, a))]
forall a b. (a -> b) -> a -> b
$ [Either Int Name] -> [(Name, a)] -> [(Either Int Name, (Name, a))]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Name, a) -> Either Int Name) -> [(Name, a)] -> [Either Int Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name -> Either Int Name
fieldish (Name -> Either Int Name)
-> ((Name, a) -> Name) -> (Name, a) -> Either Int Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, a) -> Name
forall a b. (a, b) -> a
fst) [(Name, a)]
l') [(Name, a)]
l'
  where
    l' :: [(Name, a)]
l' = Map Name a -> [(Name, a)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name a
l
    onDigit :: Maybe Int -> Char -> Maybe Int
onDigit Maybe Int
Nothing Char
_ = Maybe Int
forall a. Maybe a
Nothing
    onDigit (Just Int
d) Char
c
      | Char -> Bool
isDigit Char
c = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
10 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Char -> Int
ord Char
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Char -> Int
ord Char
'0'
      | Bool
otherwise = Maybe Int
forall a. Maybe a
Nothing
    fieldish :: Name -> Either Int Name
fieldish Name
s = Either Int Name
-> (Int -> Either Int Name) -> Maybe Int -> Either Int Name
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Name -> Either Int Name
forall a b. b -> Either a b
Right Name
s) Int -> Either Int Name
forall a b. a -> Either a b
Left (Maybe Int -> Either Int Name) -> Maybe Int -> Either Int Name
forall a b. (a -> b) -> a -> b
$ (Maybe Int -> Char -> Maybe Int) -> Maybe Int -> Text -> Maybe Int
forall a. (a -> Char -> a) -> a -> Text -> a
T.foldl' Maybe Int -> Char -> Maybe Int
onDigit (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
0) (Text -> Maybe Int) -> Text -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Name -> Text
nameToText Name
s