Skip to content

Commit d893955

Browse files
bezirgbezirg
authored
Add deriveEnum for Plinth (#7456)
1 parent 4f1c25e commit d893955

File tree

11 files changed

+454
-241
lines changed

11 files changed

+454
-241
lines changed

‎plutus-tx/changelog.d/20251203_115136_bezirg_derive_enum.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
### Added
2+
3+
- A `deriveEnum` command to derive PlutusTx.Enum instances for datatypes/newtypes, similar to Haskell's
4+
`deriving stock Enum`

‎plutus-tx/plutus-tx.cabal‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@ library
8484
PlutusTx.Data.List.TH
8585
PlutusTx.Either
8686
PlutusTx.Enum
87+
PlutusTx.Enum.Class
8788
PlutusTx.Eq
8889
PlutusTx.ErrorCodes
8990
PlutusTx.Eval
@@ -117,6 +118,7 @@ library
117118
PlutusTx.Utils
118119

119120
other-modules:
121+
PlutusTx.Enum.TH
120122
PlutusTx.IsData.Instances
121123
PlutusTx.IsData.TH
122124
PlutusTx.Lift.Instances
@@ -207,8 +209,9 @@ test-suite plutus-tx-test
207209
Blueprint.Definition.Fixture
208210
Blueprint.Definition.Spec
209211
Blueprint.Spec
210-
List.Spec
211212
Bool.Spec
213+
Enum.Spec
214+
List.Spec
212215
Rational.Laws
213216
Rational.Laws.Additive
214217
Rational.Laws.Construction
@@ -238,6 +241,7 @@ test-suite plutus-tx-test
238241
, plutus-core ^>=1.55
239242
, plutus-core:plutus-core-testlib
240243
, plutus-tx ^>=1.55
244+
, plutus-tx:plutus-tx-testlib
241245
, pretty-show
242246
, serialise
243247
, tasty

‎plutus-tx/src/PlutusTx/Enum.hs‎

Lines changed: 11 additions & 140 deletions
Original file line numberDiff line numberDiff line change
@@ -1,143 +1,14 @@
1-
{-# OPTIONS_GHC -fno-omit-interface-pragmas #-}
1+
{-# LANGUAGE OverloadedStrings #-}
2+
{-# LANGUAGE TemplateHaskell #-}
3+
{-# OPTIONS_GHC -Wno-orphans #-}
24

3-
module PlutusTx.Enum (Enum (..)) where
5+
module PlutusTx.Enum (Enum (..), deriveEnum) where
46

5-
import PlutusTx.Bool (Bool (..), otherwise)
6-
import PlutusTx.Builtins
7-
import PlutusTx.Eq ((==))
8-
import PlutusTx.ErrorCodes
9-
import PlutusTx.List
10-
import PlutusTx.Ord (Ord (..), Ordering (..))
11-
import PlutusTx.Trace
7+
import PlutusTx.Bool
8+
import PlutusTx.Enum.Class
9+
import PlutusTx.Enum.TH
10+
import PlutusTx.Ord
1211

13-
-- | Class 'Enum' defines operations on sequentially ordered types.
14-
class Enum a where
15-
{-# MINIMAL toEnum, fromEnum #-}
16-
17-
{-| The successor of a value. For numeric types, 'succ' adds 1.
18-
19-
For types that implement 'Ord', @succ x@ should be the least element
20-
that is greater than @x@, and 'error' if there is none. -}
21-
succ :: a -> a
22-
23-
{-| The predecessor of a value. For numeric types, 'pred' subtracts 1.
24-
25-
For types that implement 'Ord', @pred x@ should be the greatest element
26-
that is less than @x@, and 'error' if there is none. -}
27-
pred :: a -> a
28-
29-
-- | Convert from an 'Integer'.
30-
toEnum :: Integer -> a
31-
32-
-- | Convert to an 'Integer'.
33-
fromEnum :: a -> Integer
34-
35-
-- | Construct a list from the given range (corresponds to [a..b]).
36-
enumFromTo :: a -> a -> [a]
37-
38-
{-| Construct a list from the given range (corresponds to [a,b..c]). This
39-
has the same semantics as the Haskell version,so if a==b and c>=b then you
40-
get an infinite list, which you probably don't want in Plutus Core. -}
41-
enumFromThenTo :: a -> a -> a -> [a]
42-
43-
{-# INLINEABLE succ #-}
44-
succ x = toEnum ((`addInteger` 1) (fromEnum x))
45-
{-# INLINEABLE pred #-}
46-
pred x = toEnum ((`subtractInteger` 1) (fromEnum x))
47-
48-
{-# INLINEABLE enumFromTo #-}
49-
enumFromTo x lim = map toEnum (enumFromTo (fromEnum x) (fromEnum lim))
50-
51-
{-# INLINEABLE enumFromThenTo #-}
52-
enumFromThenTo x y lim = map toEnum (enumFromThenTo (fromEnum x) (fromEnum y) (fromEnum lim))
53-
54-
instance Enum Integer where
55-
{-# INLINEABLE succ #-}
56-
succ x = addInteger x 1
57-
58-
{-# INLINEABLE pred #-}
59-
pred x = subtractInteger x 1
60-
61-
{-# INLINEABLE toEnum #-}
62-
toEnum x = x
63-
64-
{-# INLINEABLE fromEnum #-}
65-
fromEnum x = x
66-
67-
{-# INLINEABLE enumFromTo #-}
68-
enumFromTo x lim
69-
| x > lim = []
70-
| otherwise = x : enumFromTo (succ x) lim
71-
72-
{-# INLINEABLE enumFromThenTo #-}
73-
enumFromThenTo x y lim =
74-
if delta >= 0
75-
then up_list x
76-
else dn_list x
77-
where
78-
delta = subtractInteger y x
79-
up_list x1 =
80-
if x1 > lim
81-
then []
82-
else x1 : up_list (addInteger x1 delta)
83-
dn_list x1 =
84-
if x1 < lim
85-
then []
86-
else x1 : dn_list (addInteger x1 delta)
87-
88-
instance Enum () where
89-
{-# INLINEABLE succ #-}
90-
succ _ = traceError succVoidBadArgumentError
91-
92-
{-# INLINEABLE pred #-}
93-
pred _ = traceError predVoidBadArgumentError
94-
95-
{-# INLINEABLE toEnum #-}
96-
toEnum x
97-
| x == 0 = ()
98-
| otherwise = traceError toEnumVoidBadArgumentError
99-
100-
{-# INLINEABLE fromEnum #-}
101-
fromEnum () = 0
102-
103-
instance Enum Bool where
104-
{-# INLINEABLE succ #-}
105-
succ False = True
106-
succ True = traceError succBoolBadArgumentError
107-
108-
{-# INLINEABLE pred #-}
109-
pred True = False
110-
pred False = traceError predBoolBadArgumentError
111-
112-
{-# INLINEABLE toEnum #-}
113-
toEnum n
114-
| n == 0 = False
115-
| n == 1 = True
116-
| otherwise = traceError toEnumBoolBadArgumentError
117-
118-
{-# INLINEABLE fromEnum #-}
119-
fromEnum False = 0
120-
fromEnum True = 1
121-
122-
instance Enum Ordering where
123-
{-# INLINEABLE succ #-}
124-
succ LT = EQ
125-
succ EQ = GT
126-
succ GT = traceError succOrderingBadArgumentError
127-
128-
{-# INLINEABLE pred #-}
129-
pred GT = EQ
130-
pred EQ = LT
131-
pred LT = traceError predOrderingBadArgumentError
132-
133-
{-# INLINEABLE toEnum #-}
134-
toEnum n
135-
| n == 0 = LT
136-
| n == 1 = EQ
137-
| n == 2 = GT
138-
toEnum _ = traceError toEnumOrderingBadArgumentError
139-
140-
{-# INLINEABLE fromEnum #-}
141-
fromEnum LT = 0
142-
fromEnum EQ = 1
143-
fromEnum GT = 2
12+
deriveEnum ''Bool
13+
deriveEnum ''Ordering
14+
deriveEnum ''()

‎plutus-tx/src/PlutusTx/Enum/Class.hs‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
module PlutusTx.Enum.Class (Enum (..)) where
2+
3+
import PlutusTx.Bool
4+
import PlutusTx.Builtins
5+
import PlutusTx.List
6+
import PlutusTx.Numeric
7+
import PlutusTx.Ord
8+
9+
-- | Class 'Enum' defines operations on sequentially ordered types.
10+
class Enum a where
11+
{-# MINIMAL toEnum, fromEnum #-}
12+
13+
{-| The successor of a value. For numeric types, 'succ' adds 1.
14+
15+
For types that implement 'Ord', @succ x@ should be the least element
16+
that is greater than @x@, and 'error' if there is none. -}
17+
succ :: a -> a
18+
19+
{-| The predecessor of a value. For numeric types, 'pred' subtracts 1.
20+
21+
For types that implement 'Ord', @pred x@ should be the greatest element
22+
that is less than @x@, and 'error' if there is none. -}
23+
pred :: a -> a
24+
25+
-- | Convert from an 'Integer'.
26+
toEnum :: Integer -> a
27+
28+
-- | Convert to an 'Integer'.
29+
fromEnum :: a -> Integer
30+
31+
-- | Construct a list from the given range (corresponds to [a..b]).
32+
enumFromTo :: a -> a -> [a]
33+
34+
{-| Construct a list from the given range (corresponds to [a,b..c]). This
35+
has the same semantics as the Haskell version,so if a==b and c>=b then you
36+
get an infinite list, which you probably don't want in Plutus Core. -}
37+
enumFromThenTo :: a -> a -> a -> [a]
38+
39+
{-# INLINEABLE succ #-}
40+
succ x = toEnum (fromEnum x + 1)
41+
{-# INLINEABLE pred #-}
42+
pred x = toEnum (fromEnum x - 1)
43+
44+
{-# INLINEABLE enumFromTo #-}
45+
enumFromTo x lim = map toEnum (enumFromTo (fromEnum x) (fromEnum lim))
46+
47+
{-# INLINEABLE enumFromThenTo #-}
48+
enumFromThenTo x y lim = map toEnum (enumFromThenTo (fromEnum x) (fromEnum y) (fromEnum lim))
49+
50+
instance Enum Integer where
51+
{-# INLINEABLE succ #-}
52+
succ x = addInteger x 1
53+
54+
{-# INLINEABLE pred #-}
55+
pred x = subtractInteger x 1
56+
57+
{-# INLINEABLE toEnum #-}
58+
toEnum x = x
59+
60+
{-# INLINEABLE fromEnum #-}
61+
fromEnum x = x
62+
63+
{-# INLINEABLE enumFromTo #-}
64+
enumFromTo x lim
65+
| x > lim = []
66+
| otherwise = x : enumFromTo (succ x) lim
67+
68+
{-# INLINEABLE enumFromThenTo #-}
69+
enumFromThenTo x y lim =
70+
if delta >= 0
71+
then up_list x
72+
else dn_list x
73+
where
74+
delta = subtractInteger y x
75+
up_list x1 =
76+
if x1 > lim
77+
then []
78+
else x1 : up_list (addInteger x1 delta)
79+
dn_list x1 =
80+
if x1 < lim
81+
then []
82+
else x1 : dn_list (addInteger x1 delta)

‎plutus-tx/src/PlutusTx/Enum/TH.hs‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
{-# LANGUAGE DerivingStrategies #-}
2+
{-# LANGUAGE TemplateHaskellQuotes #-}
3+
4+
module PlutusTx.Enum.TH (Enum (..), deriveEnum) where
5+
6+
import Control.Monad
7+
import Data.Deriving.Internal (varTToName)
8+
import Data.Foldable
9+
import Data.Tuple
10+
import Language.Haskell.TH as TH
11+
import Language.Haskell.TH.Datatype as TH
12+
import PlutusTx.Enum.Class
13+
import PlutusTx.ErrorCodes
14+
import PlutusTx.Trace
15+
import Prelude hiding (Bool (True), Enum (..), Eq, (&&), (==))
16+
17+
data SuccPred = Succ | Pred
18+
deriving stock (Show)
19+
20+
{-| Derive PlutusTx.Enum typeclass for datatypes, much like `deriving stock Enum` does for Haskell
21+
22+
Note: requires enabling OverloadedStrings language extension -}
23+
deriveEnum :: TH.Name -> TH.Q [TH.Dec]
24+
deriveEnum name = do
25+
TH.DatatypeInfo
26+
{ TH.datatypeName = tyConName
27+
, TH.datatypeInstTypes = tyVars0
28+
, TH.datatypeCons = cons
29+
} <-
30+
TH.reifyDatatype name
31+
let
32+
-- The purpose of the `TH.VarT . varTToName` roundtrip is to remove the kind
33+
-- signatures attached to the type variables in `tyVars0`. Otherwise, the
34+
-- `KindSignatures` extension would be needed whenever `length tyVars0 > 0`.
35+
tyVars = TH.VarT . varTToName <$> tyVars0
36+
instanceType :: TH.Type
37+
instanceType = TH.AppT (TH.ConT ''Enum) $ foldl' TH.AppT (TH.ConT tyConName) tyVars
38+
39+
table = zip (fmap IntegerL [0 ..]) (fmap constructorName cons)
40+
41+
when (null cons) $
42+
fail $
43+
"Can't make a derived instance of `Enum "
44+
++ show tyConName
45+
++ "`: "
46+
++ show tyConName
47+
++ " must must be an enumeration type (an enumeration consists of one or more nullary, non-GADT constructors)"
48+
49+
pure
50+
<$> instanceD
51+
(pure [])
52+
(pure instanceType)
53+
[ funD 'succ (fmap (deriveSuccPred Succ) (zip cons (tail (Just <$> cons) ++ [Nothing])))
54+
, TH.pragInlD 'succ TH.Inlinable TH.FunLike TH.AllPhases
55+
, funD 'pred (fmap (deriveSuccPred Pred) (zip cons (Nothing : init (Just <$> cons))))
56+
, TH.pragInlD 'pred TH.Inlinable TH.FunLike TH.AllPhases
57+
, funD 'toEnum $ fmap deriveToEnum table <> [pure toEnumDefaultClause]
58+
, TH.pragInlD 'toEnum TH.Inlinable TH.FunLike TH.AllPhases
59+
, funD 'fromEnum $ fmap (deriveFromEnum . swap) table
60+
, TH.pragInlD 'fromEnum TH.Inlinable TH.FunLike TH.AllPhases
61+
]
62+
63+
toEnumDefaultClause :: Clause
64+
toEnumDefaultClause =
65+
TH.Clause
66+
[WildP]
67+
( TH.NormalB $
68+
AppE (VarE 'traceError) (VarE 'toEnumBadArgumentError)
69+
)
70+
[]
71+
72+
deriveToEnum :: (Lit, Name) -> Q Clause
73+
deriveToEnum (l, n) = pure (TH.Clause [LitP l] (NormalB $ ConE n) [])
74+
75+
deriveFromEnum :: (Name, Lit) -> Q Clause
76+
deriveFromEnum (n, l) = pure (TH.Clause [ConP n [] []] (NormalB $ LitE l) [])
77+
78+
deriveSuccPred :: SuccPred -> (ConstructorInfo, Maybe ConstructorInfo) -> Q Clause
79+
deriveSuccPred
80+
succPred
81+
( ConstructorInfo {constructorName = nameL, constructorFields = []}
82+
, Nothing
83+
) =
84+
pure
85+
( TH.Clause
86+
[ConP nameL [] []]
87+
( NormalB $
88+
AppE
89+
(VarE 'traceError)
90+
( VarE $ case succPred of
91+
Succ -> 'succBadArgumentError
92+
Pred -> 'predBadArgumentError
93+
)
94+
)
95+
[]
96+
)
97+
deriveSuccPred
98+
_
99+
( ConstructorInfo {constructorName = nameL, constructorFields = []}
100+
, Just ConstructorInfo {constructorName = nameR, constructorFields = []}
101+
) =
102+
pure
103+
( TH.Clause
104+
[ConP nameL [] []]
105+
(NormalB $ ConE nameR)
106+
[]
107+
)
108+
deriveSuccPred _ _ = fail "Can't make a derived instance of Enum when constructor has fields"

0 commit comments

Comments
 (0)