-- \* 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' = forall a. Map Name a -> [(Name, a)]
sortFields Map Name a
fs
   in if forall (t :: * -> *). Foldable t => t Bool -> Bool
and forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall a. Eq a => a -> a -> Bool
(==) (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Name, a)]
fs') [Name]
tupleFieldNames
        then forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(Name, a)]
fs'
        else 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 = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ 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 = forall a b. (a -> b) -> [a] -> [b]
map (String -> Name
nameFromString forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map (Name -> Either Int Name
fieldish forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(Name, a)]
l') [(Name, a)]
l'
  where
    l' :: [(Name, a)]
l' = 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
_ = forall a. Maybe a
Nothing
    onDigit (Just Int
d) Char
c
      | Char -> Bool
isDigit Char
c = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Int
d forall a. Num a => a -> a -> a
* Int
10 forall a. Num a => a -> a -> a
+ Char -> Int
ord Char
c forall a. Num a => a -> a -> a
- Char -> Int
ord Char
'0'
      | Bool
otherwise = forall a. Maybe a
Nothing
    fieldish :: Name -> Either Int Name
fieldish Name
s = forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall a b. b -> Either a b
Right Name
s) forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$ forall a. (a -> Char -> a) -> a -> Text -> a
T.foldl' Maybe Int -> Char -> Maybe Int
onDigit (forall a. a -> Maybe a
Just Int
0) forall a b. (a -> b) -> a -> b
$ Name -> Text
nameToText Name
s