From 156ef8d0a79e4c5edc37d4c9dd5570050f2b72e1 Mon Sep 17 00:00:00 2001
From: crumbtoo <crumb@disroot.org>
Date: Mon, 4 Mar 2024 10:47:58 -0700
Subject: [PATCH] tysigd

---
 src/Core/Syntax.hs   | 18 ++++++++++++++++--
 src/Rlp/AltParse.y   |  7 ++++++-
 src/Rlp/AltSyntax.hs | 16 ++++++++++++++++
 src/Rlp/Lex.x        |  4 ++--
 4 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/src/Core/Syntax.hs b/src/Core/Syntax.hs
index 454118b..a12933e 100644
--- a/src/Core/Syntax.hs
+++ b/src/Core/Syntax.hs
@@ -33,6 +33,7 @@ module Core.Syntax
     , Pretty(pretty), WithTerseBinds(..)
 
     -- * Optics
+    , HasArrowSyntax(..)
     , programScDefs, programTypeSigs, programDataTags, programTyCons
     , formalising, lambdaLifting
     , HasRHS(_rhs), HasLHS(_lhs)
@@ -143,9 +144,22 @@ pattern Lit t = Fix (LitF t)
 pattern TyInt :: Type
 pattern TyInt = TyCon "Int#"
 
+class HasArrowSyntax s a b | s -> a b where
+    _arrowSyntax :: Prism' s (a, b)
+
+instance HasArrowSyntax Type Type Type where
+    _arrowSyntax = prism make unmake where
+        make (s,t) = TyFun `TyApp` s `TyApp` t
+
+        unmake (TyFun `TyApp` s `TyApp` t) = Right (s,t)
+        unmake s                           = Left s
+
 infixr 1 :->
-pattern (:->) :: Type -> Type -> Type
-pattern a :-> b = TyApp (TyApp TyFun a) b
+pattern (:->) :: HasArrowSyntax s a b
+              => a -> b -> s
+-- pattern (:->) :: Type -> Type -> Type
+pattern a :-> b <- (preview _arrowSyntax -> Just (a, b))
+    where a :-> b = _arrowSyntax # (a, b)
 
 data BindingF b a = BindingF b (ExprF b a)
     deriving (Functor, Foldable, Traversable)
diff --git a/src/Rlp/AltParse.y b/src/Rlp/AltParse.y
index 073e0e2..da00607 100644
--- a/src/Rlp/AltParse.y
+++ b/src/Rlp/AltParse.y
@@ -86,6 +86,10 @@ VS                  : ';'                   { () }
 Decl                :: { Decl PsName (RlpExpr PsName) }
                     : FunD                  { $1 }
                     | DataD                 { $1 }
+                    | TySigD                { $1 }
+
+TySigD              :: { Decl PsName (RlpExpr PsName) }
+                    : Var '::' Type         { TySigD $1 $3 }
 
 DataD               :: { Decl PsName (RlpExpr PsName) }
                     : data Con TyVars               { DataD $2 $3 [] }
@@ -104,7 +108,8 @@ Type1               :: { Type PsName }
                     | '(' Type ')'          { $2 }
 
 Type                :: { Type PsName }
-                    : AppT                  { $1 }
+                    : Type '->' Type        { $1 :-> $3 }
+                    | AppT                  { $1 }
 
 AppT                :: { Type PsName }
                     : Type1                 { $1 }
diff --git a/src/Rlp/AltSyntax.hs b/src/Rlp/AltSyntax.hs
index 16ddd42..d318bdb 100644
--- a/src/Rlp/AltSyntax.hs
+++ b/src/Rlp/AltSyntax.hs
@@ -7,6 +7,7 @@ module Rlp.AltSyntax
     , DataCon(..), Type(..)
 
     , Core.Name, PsName
+    , pattern (Core.:->)
 
     -- * Optics
     , programDecls
@@ -41,6 +42,7 @@ programDecls = lens (\ (Program ds) -> ds) (const Program)
 
 data Decl b a = FunD b [Pat b] a
               | DataD b [b] [DataCon b]
+              | TySigD b (Type b)
               deriving Show
 
 data DataCon b = DataCon b [Type b]
@@ -52,6 +54,13 @@ data Type b = VarT b
             | FunT
             deriving Show
 
+instance Core.HasArrowSyntax (Type b) (Type b) (Type b) where
+    _arrowSyntax = prism make unmake where
+        make (s,t) = FunT `AppT` s `AppT` t
+
+        unmake (FunT `AppT` s `AppT` t) = Right (s,t)
+        unmake s                        = Left s
+
 data ExprF b a = InfixEF b a a
                | LetEF Core.Rec [Binding b a] a
                | CaseEF a [Alter b a]
@@ -122,14 +131,21 @@ instance (Pretty b) => Pretty1 (Decl b) where
         cons = vcat $ zipWith (<+>) delims (pretty <$> ds)
         delims = "=" : repeat "|"
 
+    liftPrettyPrec _ _ (TySigD n t) =
+        hsep [ ttext n, ":", pretty t ]
+
 instance (Pretty b) => Pretty (DataCon b) where
     pretty (DataCon n as) = ttext n <+> hsep (prettyPrec appPrec1 <$> as)
 
+-- (->) is given prec `appPrec-1`
 instance (Pretty b) => Pretty (Type b) where
     prettyPrec _ (VarT n) = ttext n
     prettyPrec _ (ConT n) = ttext n
+    prettyPrec p (s Core.:-> t) = maybeParens (p>appPrec-1) $
+        hsep [ prettyPrec appPrec s, "->", prettyPrec (appPrec-1) t ]
     prettyPrec p (AppT f x) = maybeParens (p>appPrec) $
         prettyPrec appPrec f <+> prettyPrec appPrec1 x
+    prettyPrec p FunT = maybeParens (p>0) "->"
 
 instance (Pretty b) => Pretty (Pat b) where
     prettyPrec p (VarP b) = prettyPrec p b
diff --git a/src/Rlp/Lex.x b/src/Rlp/Lex.x
index 93cac61..88380e3 100644
--- a/src/Rlp/Lex.x
+++ b/src/Rlp/Lex.x
@@ -61,7 +61,7 @@ $asciisym       = [\!\#\$\%\&\*\+\.\/\<\=\>\?\@\\\^\|\-\~\:]
     |infixr|infixl|infix
 
 @reservedop =
-    "=" | \\ | "->" | "|" | "::"
+    "=" | \\ | "->" | "|" | ":"
 
 rlp :-
     
@@ -167,7 +167,7 @@ lexReservedName = \case
 lexReservedOp :: Text -> RlpToken
 lexReservedOp = \case
     "="     -> TokenEquals
-    "::"    -> TokenHasType
+    ":"     -> TokenHasType
     "|"     -> TokenPipe
     "->"    -> TokenArrow
     s       -> error (show s)