-- Copyright 2021 Google LLC
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--      http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

module MLIR.AST.PatternUtil
  ( pattern I32ArrayAttr
  , pattern I64ArrayAttr
  , pattern AffineMapArrayAttr
  , DummyIx
  ) where

import Data.Traversable
import Data.Array

import MLIR.AST
import qualified MLIR.AST.Dialect.Affine as Affine

unwrapI32ArrayAttr :: Attribute -> Maybe [Int]
unwrapI32ArrayAttr :: Attribute -> Maybe [Int]
unwrapI32ArrayAttr (ArrayAttr [Attribute]
vals) = [Attribute] -> (Attribute -> Maybe Int) -> Maybe [Int]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Attribute]
vals \case
  IntegerAttr (IntegerType Signedness
Signed UInt
32) Int
v -> Int -> Maybe Int
forall a. a -> Maybe a
Just Int
v
  Attribute
_                                     -> Maybe Int
forall a. Maybe a
Nothing
unwrapI32ArrayAttr Attribute
_ = Maybe [Int]
forall a. Maybe a
Nothing

pattern I32ArrayAttr :: [Int] -> Attribute
pattern $bI32ArrayAttr :: [Int] -> Attribute
$mI32ArrayAttr :: forall r. Attribute -> ([Int] -> r) -> (Void# -> r) -> r
I32ArrayAttr vals <- (unwrapI32ArrayAttr -> Just vals)
  where I32ArrayAttr [Int]
vals = [Attribute] -> Attribute
ArrayAttr ([Attribute] -> Attribute) -> [Attribute] -> Attribute
forall a b. (a -> b) -> a -> b
$ (Int -> Attribute) -> [Int] -> [Attribute]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Type -> Int -> Attribute
IntegerAttr (Signedness -> UInt -> Type
IntegerType Signedness
Signed UInt
32)) [Int]
vals

unwrapI64ArrayAttr :: Attribute -> Maybe [Int]
unwrapI64ArrayAttr :: Attribute -> Maybe [Int]
unwrapI64ArrayAttr (ArrayAttr [Attribute]
vals) = [Attribute] -> (Attribute -> Maybe Int) -> Maybe [Int]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Attribute]
vals \case
  IntegerAttr (IntegerType Signedness
Signed UInt
64) Int
v -> Int -> Maybe Int
forall a. a -> Maybe a
Just Int
v
  Attribute
_                                     -> Maybe Int
forall a. Maybe a
Nothing
unwrapI64ArrayAttr Attribute
_ = Maybe [Int]
forall a. Maybe a
Nothing

pattern I64ArrayAttr :: [Int] -> Attribute
pattern $bI64ArrayAttr :: [Int] -> Attribute
$mI64ArrayAttr :: forall r. Attribute -> ([Int] -> r) -> (Void# -> r) -> r
I64ArrayAttr vals <- (unwrapI64ArrayAttr -> Just vals)
  where I64ArrayAttr [Int]
vals = [Attribute] -> Attribute
ArrayAttr ([Attribute] -> Attribute) -> [Attribute] -> Attribute
forall a b. (a -> b) -> a -> b
$ (Int -> Attribute) -> [Int] -> [Attribute]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Type -> Int -> Attribute
IntegerAttr (Signedness -> UInt -> Type
IntegerType Signedness
Signed UInt
64)) [Int]
vals

unwrapAffineMapArrayAttr :: Attribute -> Maybe [Affine.Map]
unwrapAffineMapArrayAttr :: Attribute -> Maybe [Map]
unwrapAffineMapArrayAttr (ArrayAttr [Attribute]
vals) = [Attribute] -> (Attribute -> Maybe Map) -> Maybe [Map]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Attribute]
vals \case
  AffineMapAttr Map
m -> Map -> Maybe Map
forall a. a -> Maybe a
Just Map
m
  Attribute
_               -> Maybe Map
forall a. Maybe a
Nothing
unwrapAffineMapArrayAttr Attribute
_ = Maybe [Map]
forall a. Maybe a
Nothing

pattern AffineMapArrayAttr :: [Affine.Map] -> Attribute
pattern $bAffineMapArrayAttr :: [Map] -> Attribute
$mAffineMapArrayAttr :: forall r. Attribute -> ([Map] -> r) -> (Void# -> r) -> r
AffineMapArrayAttr vals <- (unwrapAffineMapArrayAttr -> Just vals)
  where AffineMapArrayAttr [Map]
vals = [Attribute] -> Attribute
ArrayAttr ([Attribute] -> Attribute) -> [Attribute] -> Attribute
forall a b. (a -> b) -> a -> b
$ (Map -> Attribute) -> [Map] -> [Attribute]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Map -> Attribute
AffineMapAttr [Map]
vals

data DummyIx
deriving instance Eq DummyIx
deriving instance Ord DummyIx
deriving instance Show DummyIx
instance Ix DummyIx where
  range :: (DummyIx, DummyIx) -> [DummyIx]
range   (DummyIx, DummyIx)
_   = String -> [DummyIx]
forall a. HasCallStack => String -> a
error String
"Invalid index"
  index :: (DummyIx, DummyIx) -> DummyIx -> Int
index   (DummyIx, DummyIx)
_ DummyIx
_ = String -> Int
forall a. HasCallStack => String -> a
error String
"Invalid index"
  inRange :: (DummyIx, DummyIx) -> DummyIx -> Bool
inRange (DummyIx, DummyIx)
_ DummyIx
_ = String -> Bool
forall a. HasCallStack => String -> a
error String
"Invalid index"