simplify multiwya if

src/Juvix/Compiler/Concrete/Language.hs

@@ -59,10 +59,15 @@ type family FieldArgIxType s = res | res -> s where
   FieldArgIxType 'Parsed = ()
   FieldArgIxType 'Scoped = Int
 
+type SideIfBranchConditionType :: Stage -> IfBranchKind -> GHC.Type
+type family SideIfBranchConditionType s k = res where
+  SideIfBranchConditionType s 'BranchIfBool = ExpressionType s
+  SideIfBranchConditionType _ 'BranchIfElse = ()
+
 type IfBranchConditionType :: Stage -> IfBranchKind -> GHC.Type
 type family IfBranchConditionType s k = res where
   IfBranchConditionType s 'BranchIfBool = ExpressionType s
-  IfBranchConditionType _ 'BranchIfElse = ()
+  IfBranchConditionType _ 'BranchIfElse = Irrelevant KeywordRef
 
 type ModuleIdType :: Stage -> ModuleIsTop -> GHC.Type
 type family ModuleIdType s t = res where
@@ -1712,7 +1717,7 @@ deriving stock instance Ord (Let 'Scoped)
 data SideIfBranch (s :: Stage) (k :: IfBranchKind) = SideIfBranch
   { _sideIfBranchPipe :: Irrelevant (Maybe KeywordRef),
     _sideIfBranchKw :: Irrelevant KeywordRef,
-    _sideIfBranchCondition :: IfBranchConditionType s k,
+    _sideIfBranchCondition :: SideIfBranchConditionType s k,
     _sideIfBranchAssignKw :: Irrelevant KeywordRef,
     _sideIfBranchBody :: ExpressionType s
   }
@@ -1946,66 +1951,58 @@ deriving stock instance Ord (NewCase 'Parsed)
 
 deriving stock instance Ord (NewCase 'Scoped)
 
-data IfBranch (s :: Stage) = IfBranch
+data IfBranch (s :: Stage) (k :: IfBranchKind) = IfBranch
   { _ifBranchPipe :: Irrelevant KeywordRef,
     _ifBranchAssignKw :: Irrelevant KeywordRef,
-    _ifBranchCondition :: ExpressionType s,
+    _ifBranchCondition :: IfBranchConditionType s k,
     _ifBranchExpression :: ExpressionType s
   }
   deriving stock (Generic)
 
-instance Serialize (IfBranch 'Scoped)
+instance Serialize (IfBranch 'Scoped 'BranchIfBool)
 
-instance NFData (IfBranch 'Scoped)
+instance Serialize (IfBranch 'Scoped 'BranchIfElse)
 
-instance Serialize (IfBranch 'Parsed)
+instance NFData (IfBranch 'Scoped 'BranchIfBool)
 
-instance NFData (IfBranch 'Parsed)
+instance NFData (IfBranch 'Scoped 'BranchIfElse)
 
-deriving stock instance Show (IfBranch 'Parsed)
+instance Serialize (IfBranch 'Parsed 'BranchIfBool)
 
-deriving stock instance Show (IfBranch 'Scoped)
+instance Serialize (IfBranch 'Parsed 'BranchIfElse)
 
-deriving stock instance Eq (IfBranch 'Parsed)
+instance NFData (IfBranch 'Parsed 'BranchIfElse)
 
-deriving stock instance Eq (IfBranch 'Scoped)
+instance NFData (IfBranch 'Parsed 'BranchIfBool)
 
-deriving stock instance Ord (IfBranch 'Parsed)
+deriving stock instance Show (IfBranch 'Parsed 'BranchIfElse)
 
-deriving stock instance Ord (IfBranch 'Scoped)
+deriving stock instance Show (IfBranch 'Parsed 'BranchIfBool)
 
-data IfBranchElse (s :: Stage) = IfBranchElse
-  { _ifBranchElsePipe :: Irrelevant KeywordRef,
-    _ifBranchElseAssignKw :: Irrelevant KeywordRef,
-    _ifBranchElseKw :: Irrelevant KeywordRef,
-    _ifBranchElseExpression :: ExpressionType s
-  }
-  deriving stock (Generic)
+deriving stock instance Show (IfBranch 'Scoped 'BranchIfElse)
 
-instance Serialize (IfBranchElse 'Scoped)
+deriving stock instance Show (IfBranch 'Scoped 'BranchIfBool)
 
-instance NFData (IfBranchElse 'Scoped)
+deriving stock instance Eq (IfBranch 'Parsed 'BranchIfElse)
 
-instance Serialize (IfBranchElse 'Parsed)
+deriving stock instance Eq (IfBranch 'Parsed 'BranchIfBool)
 
-instance NFData (IfBranchElse 'Parsed)
+deriving stock instance Eq (IfBranch 'Scoped 'BranchIfElse)
 
-deriving stock instance Show (IfBranchElse 'Parsed)
+deriving stock instance Eq (IfBranch 'Scoped 'BranchIfBool)
 
-deriving stock instance Show (IfBranchElse 'Scoped)
+deriving stock instance Ord (IfBranch 'Parsed 'BranchIfElse)
 
-deriving stock instance Eq (IfBranchElse 'Parsed)
+deriving stock instance Ord (IfBranch 'Parsed 'BranchIfBool)
 
-deriving stock instance Eq (IfBranchElse 'Scoped)
+deriving stock instance Ord (IfBranch 'Scoped 'BranchIfElse)
 
-deriving stock instance Ord (IfBranchElse 'Parsed)
-
-deriving stock instance Ord (IfBranchElse 'Scoped)
+deriving stock instance Ord (IfBranch 'Scoped 'BranchIfBool)
 
 data If (s :: Stage) = If
   { _ifKw :: KeywordRef,
-    _ifBranches :: NonEmpty (IfBranch s),
-    _ifBranchElse :: IfBranchElse s
+    _ifBranches :: [IfBranch s 'BranchIfBool],
+    _ifBranchElse :: IfBranch s 'BranchIfElse
   }
   deriving stock (Generic)
 
@@ -2657,7 +2654,6 @@ makeLenses ''Case
 makeLenses ''CaseBranch
 makeLenses ''If
 makeLenses ''IfBranch
-makeLenses ''IfBranchElse
 makeLenses ''PatternBinding
 makeLenses ''PatternAtoms
 makeLenses ''ExpressionAtoms
@@ -2891,12 +2887,9 @@ instance (SingI s) => HasLoc (CaseBranch s) where
       branchLoc :: Interval
       branchLoc = getLoc (c ^. caseBranchRhs)
 
-instance (SingI s) => HasLoc (IfBranch s) where
+instance (SingI s) => HasLoc (IfBranch s k) where
   getLoc c = getLoc (c ^. ifBranchPipe) <> getLocExpressionType (c ^. ifBranchExpression)
 
-instance (SingI s) => HasLoc (IfBranchElse s) where
-  getLoc c = getLoc (c ^. ifBranchElsePipe) <> getLocExpressionType (c ^. ifBranchElseExpression)
-
 instance (SingI s) => HasLoc (Case s) where
   getLoc c = getLoc (c ^. caseKw) <> getLoc (c ^. caseBranches . to last)
 

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

@@ -616,32 +616,37 @@ ppCase isTop Case {..} = do
               Just p -> ppCode p
               Nothing -> ppCode Kw.kwPipe
 
-instance (SingI s) => PrettyPrint (IfBranch s) where
+instance (SingI s) => PrettyPrint (IfBranch s 'BranchIfBool) where
   ppCode IfBranch {..} = do
-    let cond' = ppExpressionType _ifBranchCondition
+    let pipe' = ppCode _ifBranchPipe
+        cond' = ppExpressionType _ifBranchCondition
         e' = ppExpressionType _ifBranchExpression
-    cond' <+> ppCode _ifBranchAssignKw <> oneLineOrNext e'
+    pipe' <+> cond' <+> ppCode _ifBranchAssignKw <> oneLineOrNext e'
 
-ppIfBranchElse :: forall r s. (Members '[ExactPrint, Reader Options] r, SingI s) => IsTop -> IfBranchElse s -> Sem r ()
-ppIfBranchElse isTop IfBranchElse {..} = do
-  let e' = ppMaybeTopExpression isTop _ifBranchElseExpression
-  ppCode _ifBranchElseKw <+> ppCode _ifBranchElseAssignKw <> oneLineOrNext e'
+ppIfBranchElse ::
+  forall r s.
+  (Members '[ExactPrint, Reader Options] r, SingI s) =>
+  IsTop ->
+  IfBranch s 'BranchIfElse ->
+  Sem r ()
+ppIfBranchElse isTop IfBranch {..} = do
+  let e' = ppMaybeTopExpression isTop _ifBranchExpression
+  ppCode _ifBranchCondition <+> ppCode _ifBranchAssignKw <> oneLineOrNext e'
 
 ppIf :: forall r s. (Members '[ExactPrint, Reader Options] r, SingI s) => IsTop -> If s -> Sem r ()
 ppIf isTop If {..} = do
-  ppCode _ifKw <+> hardline <> indent (vsepHard (ppIfBranch <$> _ifBranches) <> hardline <> ppIfBranchElse' _ifBranchElse)
+  ppCode _ifKw
+    <+> hardline
+      <> indent
+        ( vsepHard (ppIfBranch <$> _ifBranches)
+            <> hardline
+            <> ppIfBranch _ifBranchElse
+        )
   where
-    ppIfBranch :: IfBranch s -> Sem r ()
-    ppIfBranch b = pipeHelper <+> ppCode b
-      where
-        pipeHelper :: Sem r ()
-        pipeHelper = ppCode (b ^. ifBranchPipe . unIrrelevant)
-
-    ppIfBranchElse' :: IfBranchElse s -> Sem r ()
-    ppIfBranchElse' b = pipeHelper <+> ppIfBranchElse isTop b
-      where
-        pipeHelper :: Sem r ()
-        pipeHelper = ppCode (b ^. ifBranchElsePipe . unIrrelevant)
+    ppIfBranch :: forall k. (SingI k) => IfBranch s k -> Sem r ()
+    ppIfBranch b = case sing :: SIfBranchKind k of
+      SBranchIfBool -> ppCode b
+      SBranchIfElse -> ppIfBranchElse isTop b
 
 instance PrettyPrint Universe where
   ppCode Universe {..} = ppCode _universeKw <+?> (noLoc . pretty <$> _universeLevel)

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

@@ -2193,7 +2193,7 @@ checkRhsExpression ::
          InfoTableBuilder,
          Reader InfoTable,
          NameIdGen,
-         Reader EntryPoint
+         Reader Package
        ]
       r
   ) =>
@@ -2219,7 +2219,7 @@ checkSideIfBranch ::
          InfoTableBuilder,
          Reader InfoTable,
          NameIdGen,
-         Reader EntryPoint
+         Reader Package
        ]
       r
   ) =>
@@ -2250,7 +2250,7 @@ checkSideIfs ::
          InfoTableBuilder,
          Reader InfoTable,
          NameIdGen,
-         Reader EntryPoint
+         Reader Package
        ]
       r
   ) =>
@@ -2267,7 +2267,7 @@ checkSideIfs SideIfs {..} = do
 
 checkCaseBranchRhs ::
   forall r.
-  (Members '[HighlightBuilder, Reader ScopeParameters, Error ScoperError, State Scope, State ScoperState, InfoTableBuilder, Reader InfoTable, NameIdGen, Reader EntryPoint] r) =>
+  (Members '[HighlightBuilder, Reader ScopeParameters, Error ScoperError, State Scope, State ScoperState, InfoTableBuilder, Reader InfoTable, NameIdGen, Reader Package] r) =>
   CaseBranchRhs 'Parsed ->
   Sem r (CaseBranchRhs 'Scoped)
 checkCaseBranchRhs = \case
@@ -2305,12 +2305,14 @@ checkCase Case {..} = do
       }
 
 checkIfBranch ::
-  forall r.
-  (Members '[HighlightBuilder, Reader ScopeParameters, Error ScoperError, State Scope, State ScoperState, InfoTableBuilder, Reader InfoTable, NameIdGen, Reader Package] r) =>
-  IfBranch 'Parsed ->
-  Sem r (IfBranch 'Scoped)
+  forall r k.
+  (SingI k, Members '[HighlightBuilder, Reader ScopeParameters, Error ScoperError, State Scope, State ScoperState, InfoTableBuilder, Reader InfoTable, NameIdGen, Reader Package] r) =>
+  IfBranch 'Parsed k ->
+  Sem r (IfBranch 'Scoped k)
 checkIfBranch IfBranch {..} = withLocalScope $ do
-  cond' <- checkParseExpressionAtoms _ifBranchCondition
+  cond' <- case sing :: SIfBranchKind k of
+    SBranchIfBool -> checkParseExpressionAtoms _ifBranchCondition
+    SBranchIfElse -> return _ifBranchCondition
   expression' <- checkParseExpressionAtoms _ifBranchExpression
   return $
     IfBranch
@@ -2319,26 +2321,13 @@ checkIfBranch IfBranch {..} = withLocalScope $ do
         ..
       }
 
-checkIfBranchElse ::
-  forall r.
-  (Members '[HighlightBuilder, Reader ScopeParameters, Error ScoperError, State Scope, State ScoperState, InfoTableBuilder, Reader InfoTable, NameIdGen, Reader Package] r) =>
-  IfBranchElse 'Parsed ->
-  Sem r (IfBranchElse 'Scoped)
-checkIfBranchElse IfBranchElse {..} = withLocalScope $ do
-  expression' <- checkParseExpressionAtoms _ifBranchElseExpression
-  return $
-    IfBranchElse
-      { _ifBranchElseExpression = expression',
-        ..
-      }
-
 checkIf ::
   (Members '[HighlightBuilder, Reader ScopeParameters, Error ScoperError, State Scope, State ScoperState, InfoTableBuilder, Reader InfoTable, NameIdGen, Reader Package] r) =>
   If 'Parsed ->
   Sem r (If 'Scoped)
 checkIf If {..} = do
   ifBranches' <- mapM checkIfBranch _ifBranches
-  ifBranchElse' <- checkIfBranchElse _ifBranchElse
+  ifBranchElse' <- checkIfBranch _ifBranchElse
   return $
     If
       { _ifBranchElse = ifBranchElse',

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

@@ -825,7 +825,7 @@ expressionAtom =
       <|> AtomNamedApplicationNew <$> namedApplicationNew
       <|> AtomNamedApplication <$> namedApplication
       <|> AtomList <$> parseList
-      <|> either AtomIf AtomIdentifier <$> multiwayIf
+      <|> AtomIf <$> multiwayIf
       <|> AtomIdentifier <$> name
       <|> AtomUniverse <$> universe
       <|> AtomLambda <$> lambda
@@ -1179,40 +1179,30 @@ case_ = P.label "case" $ do
   _caseBranches <- braces (pipeSep1 caseBranch) <|> pipeSep1 caseBranch
   return Case {..}
 
-ifBranch' :: (Members '[ParserResultBuilder, PragmasStash, JudocStash] r) => Irrelevant KeywordRef -> ParsecS r (IfBranch 'Parsed)
-ifBranch' _ifBranchPipe = do
-  _ifBranchCondition <- parseExpressionAtoms
+ifBranch ::
+  forall r k.
+  (SingI k, Members '[ParserResultBuilder, PragmasStash, JudocStash] r) =>
+  ParsecS r (IfBranch 'Parsed k)
+ifBranch = do
+  _ifBranchPipe <- Irrelevant <$> pipeHelper
+  _ifBranchCondition <- case sing :: SIfBranchKind k of
+    SBranchIfBool -> parseExpressionAtoms
+    SBranchIfElse -> Irrelevant <$> kw kwElse
   _ifBranchAssignKw <- Irrelevant <$> kw kwAssign
   _ifBranchExpression <- parseExpressionAtoms
   return IfBranch {..}
-
-parseIfBranchElse' :: (Members '[ParserResultBuilder, PragmasStash, JudocStash] r) => Irrelevant KeywordRef -> ParsecS r (IfBranchElse 'Parsed)
-parseIfBranchElse' _ifBranchElsePipe = do
-  _ifBranchElseKw <- Irrelevant <$> kw kwElse
-  _ifBranchElseAssignKw <- Irrelevant <$> kw kwAssign
-  _ifBranchElseExpression <- parseExpressionAtoms
-  return IfBranchElse {..}
-
-mumultiwayIfBranch' :: (Members '[ParserResultBuilder, PragmasStash, JudocStash] r) => KeywordRef -> Irrelevant KeywordRef -> [IfBranch 'Parsed] -> ParsecS r (If 'Parsed)
-
-multiwayIfBranch' _ifKw pipeKw brs = do
-  br <- ifBranch' pipeKw
-  multiwayIf' _ifKw (br : brs)
-
-multiwayIfBranchElse' :: (Members '[ParserResultBuilder, PragmasStash, JudocStash] r) => KeywordRef -> Irrelevant KeywordRef -> [IfBranch 'Parsed] -> ParsecS r (If 'Parsed)
-multiwayIfBranchElse' _ifKw pipeKw brs = do
-  off <- P.getOffset
-  _ifBranchElse <- parseIfBranchElse' pipeKw
-  case nonEmpty (reverse brs) of
-    Nothing -> parseFailure off "A multiway if must have at least one condition branch"
-    Just _ifBranches -> return If {..}
+  where
+    pipeHelper :: ParsecS r KeywordRef
+    pipeHelper = case sing :: SIfBranchKind k of
+      SBranchIfBool -> P.try (kw kwPipe <* P.notFollowedBy (kw kwElse))
+      SBranchIfElse -> kw kwElse
 
 multiwayIf :: (Members '[ParserResultBuilder, PragmasStash, JudocStash] r) => ParsecS r (If 'Parsed)
 multiwayIf = do
   _ifKw <- kw kwIf
-  multiwayIf' _ifKw []
-  pipeKw <- Irrelevant <$> kw kwPipe
-  multiwayIfBranchElse' _ifKw pipeKw brs <|> multiwayIfBranch' _ifKw pipeKw brs
+  _ifBranches <- many ifBranch
+  _ifBranchElse <- ifBranch
+  return If {..}
 
 --------------------------------------------------------------------------------
 -- Universe expression

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

@@ -942,12 +942,12 @@ goExpression = \case
     goIf :: Concrete.If 'Scoped -> Sem r Internal.Expression
     goIf e@Concrete.If {..} = do
       if_ <- getBuiltinName (getLoc e) BuiltinBoolIf
-      go if_ (toList _ifBranches)
+      go if_ _ifBranches
       where
-        go :: Internal.Name -> [Concrete.IfBranch 'Scoped] -> Sem r Internal.Expression
+        go :: Internal.Name -> [Concrete.IfBranch 'Scoped 'BranchIfBool] -> Sem r Internal.Expression
         go if_ = \case
           [] ->
-            goExpression (_ifBranchElse ^. Concrete.ifBranchElseExpression)
+            goExpression (_ifBranchElse ^. Concrete.ifBranchExpression)
           Concrete.IfBranch {..} : brs -> do
             c <- goExpression _ifBranchCondition
             b1 <- goExpression _ifBranchExpression
@@ -1083,7 +1083,7 @@ gRhsExpression RhsExpression {..} = goExpression _rhsExpression
 goSideIfBranch ::
   forall r.
   (Members '[Reader DefaultArgsStack, Builtins, NameIdGen, Error ScoperError, Reader Pragmas, Reader S.InfoTable] r) =>
-  BranchIfBranch 'Scoped 'BranchIfBool ->
+  SideIfBranch 'Scoped 'BranchIfBool ->
   Sem r Internal.SideIfBranch
 goSideIfBranch s = do
   cond' <- goExpression (s ^. sideIfBranchCondition)