The formatter respects the ascii function arrow

src/Juvix/Compiler/Abstract/Translation/FromConcrete.hs

@@ -446,9 +446,9 @@ goUniverse :: Universe -> Universe
 goUniverse = id
 
 goFunction :: (Members '[Error ScoperError, InfoTableBuilder] r) => Function 'Scoped -> Sem r Abstract.Function
-goFunction (Function l r) = do
-  params <- goFunctionParameters l
-  ret <- goExpression r
+goFunction f = do
+  params <- goFunctionParameters (f ^. funParameters)
+  ret <- goExpression (f ^. funReturn)
   return $
     Abstract.Function (head params) $
       foldr (\param acc -> Abstract.ExpressionFunction $ Abstract.Function param acc) ret (NonEmpty.tail params)

src/Juvix/Compiler/Concrete/Data/ParsedInfoTableBuilder.hs

@@ -18,7 +18,6 @@ import Juvix.Compiler.Concrete.Data.Literal
 import Juvix.Compiler.Concrete.Data.ParsedInfoTable
 import Juvix.Compiler.Concrete.Data.ParsedItem
 import Juvix.Compiler.Concrete.Language
-import Juvix.Data.Keyword
 import Juvix.Prelude
 
 data InfoTableBuilder m a where

src/Juvix/Compiler/Concrete/Language.hs

@@ -4,6 +4,7 @@ module Juvix.Compiler.Concrete.Language
   ( module Juvix.Compiler.Concrete.Language,
     module Juvix.Compiler.Concrete.Data.Name,
     module Juvix.Compiler.Concrete.Data.NameRef,
+    module Juvix.Data.Keyword,
     module Juvix.Compiler.Concrete.Data.Builtins,
     module Juvix.Compiler.Concrete.Data.Literal,
     module Juvix.Data,
@@ -617,6 +618,7 @@ deriving stock instance
 -- | Function *type* representation
 data Function (s :: Stage) = Function
   { _funParameters :: FunctionParameters s,
+    _funKw :: KeywordRef,
     _funReturn :: ExpressionType s
   }
 
@@ -865,7 +867,7 @@ data ExpressionAtom (s :: Stage)
   | AtomLetBlock (LetBlock s)
   | AtomUniverse Universe
   | AtomFunction (Function s)
-  | AtomFunArrow
+  | AtomFunArrow KeywordRef
   | AtomLiteral LiteralLoc
   | AtomParens (ExpressionType s)
 
@@ -1259,8 +1261,13 @@ instance HasAtomicity (ExpressionAtoms 'Parsed) where
   atomicity ExpressionAtoms {..} = case _expressionAtoms of
     (_ :| []) -> Atom
     (_ :| _)
-      | AtomFunArrow `elem` _expressionAtoms -> Aggregate funFixity
+      | any isArrow _expressionAtoms -> Aggregate funFixity
       | otherwise -> Aggregate appFixity
+      where
+        isArrow :: ExpressionAtom s -> Bool
+        isArrow = \case
+          AtomFunArrow {} -> True
+          _ -> False
 
 instance
   ( Eq (ExpressionType s),

src/Juvix/Compiler/Concrete/Pretty/Base.hs

@@ -14,7 +14,6 @@ import Juvix.Compiler.Concrete.Language
 import Juvix.Compiler.Concrete.Pretty.Options
 import Juvix.Data.Ape
 import Juvix.Data.CodeAnn
-import Juvix.Data.Keyword
 import Juvix.Extra.Strings qualified as Str
 import Juvix.Prelude
 import Juvix.Prelude.Pretty qualified as PP
@@ -450,7 +449,8 @@ instance (SingI s) => PrettyCode (Function s) where
   ppCode Function {..} = do
     funParameter' <- ppCode _funParameters
     funReturn' <- ppRightExpression' funFixity _funReturn
-    return $ funParameter' <+> kwArrowR <+> funReturn'
+    funKw' <- ppCode _funKw
+    return $ funParameter' <+> funKw' <+> funReturn'
     where
       ppRightExpression' = case sing :: SStage s of
         SParsed -> ppRightExpression
@@ -774,7 +774,7 @@ instance (SingI s) => PrettyCode (ExpressionAtom s) where
     AtomUniverse uni -> ppCode uni
     AtomFunction fun -> ppCode fun
     AtomLiteral lit -> ppCode lit
-    AtomFunArrow -> return kwArrowR
+    AtomFunArrow a -> ppCode a
     AtomParens e -> parens <$> ppExpression e
     AtomBraces e -> braces <$> ppExpression (e ^. withLocParam)
     AtomHole w -> ppHole w

src/Juvix/Compiler/Concrete/Translation/FromParsed/Analysis/Scoping.hs

@@ -940,7 +940,8 @@ checkFunction Function {..} = do
   return
     Function
       { _funParameters = funParameters',
-        _funReturn = funReturn'
+        _funReturn = funReturn',
+        _funKw
       }
   where
     go :: Maybe (SymbolType 'Scoped) -> (Sem r a -> Sem r a) -> (Sem r a -> Sem r a)
@@ -1220,7 +1221,7 @@ checkExpressionAtom e = case e of
   AtomFunction fun -> AtomFunction <$> checkFunction fun
   AtomParens par -> AtomParens <$> checkParens par
   AtomBraces br -> AtomBraces <$> traverseOf withLocParam checkParseExpressionAtoms br
-  AtomFunArrow -> return AtomFunArrow
+  AtomFunArrow a -> return (AtomFunArrow a)
   AtomHole h -> AtomHole <$> checkHole h
   AtomLiteral l -> return (AtomLiteral l)
 
@@ -1377,14 +1378,19 @@ makeExpressionTable2 (ExpressionAtoms atoms _) = [appOpExplicit] : operators ++
 
     -- Non-dependent function type: A → B
     functionOp :: P.Operator Parse Expression
-    functionOp = P.InfixR (nonDepFun <$ P.single AtomFunArrow)
+    functionOp = P.InfixR (nonDepFun <$> P.token getArrow mempty)
       where
-        nonDepFun :: Expression -> Expression -> Expression
-        nonDepFun a b =
+        getArrow :: ExpressionAtom 'Scoped -> Maybe KeywordRef
+        getArrow = \case
+          AtomFunArrow r -> return r
+          _ -> Nothing
+        nonDepFun :: KeywordRef -> Expression -> Expression -> Expression
+        nonDepFun _funKw a b =
           ExpressionFunction
             Function
               { _funParameters = param,
-                _funReturn = b
+                _funReturn = b,
+                _funKw
               }
           where
             param =

src/Juvix/Compiler/Concrete/Translation/FromSource.hs

@@ -367,7 +367,7 @@ expressionAtom =
       <|> (AtomCase <$> case_)
       <|> (AtomFunction <$> function)
       <|> (AtomLetBlock <$> letBlock)
-      <|> (AtomFunArrow <$ kw kwRightArrow)
+      <|> (AtomFunArrow <$> kw kwRightArrow)
       <|> (AtomHole <$> hole)
       <|> parens (AtomParens <$> parseExpressionAtoms)
       <|> braces (AtomBraces <$> withLoc parseExpressionAtoms)
@@ -522,7 +522,7 @@ functionParams = do
 function :: (Members '[InfoTableBuilder, JudocStash, NameIdGen] r) => ParsecS r (Function 'Parsed)
 function = do
   _funParameters <- functionParams
-  kw kwRightArrow
+  _funKw <- kw kwRightArrow
   _funReturn <- parseExpressionAtoms
   return Function {..}
 
@@ -668,7 +668,7 @@ atomicExpression :: (Members '[InfoTableBuilder, JudocStash, NameIdGen] r) => Pa
 atomicExpression = do
   (atom, loc) <- interval expressionAtom
   case atom of
-    AtomFunArrow -> P.failure Nothing mempty
+    AtomFunArrow {} -> P.failure Nothing mempty
     _ -> return ()
   return $ ExpressionAtoms (NonEmpty.singleton atom) loc