Soft casting with if/is

crates/aiken-lang/src/ast.rs

@@ -1679,13 +1679,15 @@ pub type TypedAssignmentKind = AssignmentKind<()>;
 
 #[derive(Debug, Clone, PartialEq, Eq, Copy, serde::Serialize, serde::Deserialize)]
 pub enum AssignmentKind<T> {
+    Is,
     Let { backpassing: T },
     Expect { backpassing: T },
 }
 
 impl From<UntypedAssignmentKind> for TypedAssignmentKind {
     fn from(kind: UntypedAssignmentKind) -> TypedAssignmentKind {
         match kind {
+            AssignmentKind::Is => AssignmentKind::Is,
             AssignmentKind::Let { .. } => AssignmentKind::Let { backpassing: () },
             AssignmentKind::Expect { .. } => AssignmentKind::Expect { backpassing: () },
         }
@@ -1701,8 +1703,13 @@ impl<T> AssignmentKind<T> {
         matches!(self, AssignmentKind::Expect { .. })
     }
 
+    pub fn if_is(&self) -> bool {
+        matches!(self, AssignmentKind::Is)
+    }
+
     pub fn location_offset(&self) -> usize {
         match self {
+            AssignmentKind::Is => 2,
             AssignmentKind::Let { .. } => 3,
             AssignmentKind::Expect { .. } => 6,
         }
@@ -1712,6 +1719,7 @@ impl<T> AssignmentKind<T> {
 impl AssignmentKind<bool> {
     pub fn is_backpassing(&self) -> bool {
         match self {
+            Self::Is => unreachable!(),
             Self::Let { backpassing } | Self::Expect { backpassing } => *backpassing,
         }
     }
@@ -1724,6 +1732,10 @@ impl<T: Default> AssignmentKind<T> {
         }
     }
 
+    pub fn is() -> Self {
+        AssignmentKind::Is
+    }
+
     pub fn expect() -> Self {
         AssignmentKind::Expect {
             backpassing: Default::default(),
@@ -1885,13 +1897,14 @@ impl TypedClauseGuard {
     }
 }
 
-pub type TypedIfBranch = IfBranch<TypedExpr>;
-pub type UntypedIfBranch = IfBranch<UntypedExpr>;
+pub type TypedIfBranch = IfBranch<TypedExpr, TypedPattern>;
+pub type UntypedIfBranch = IfBranch<UntypedExpr, AssignmentPattern>;
 
 #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
-pub struct IfBranch<Expr> {
+pub struct IfBranch<Expr, Is> {
     pub condition: Expr,
     pub body: Expr,
+    pub is: Option<Is>,
     pub location: Span,
 }
 

crates/aiken-lang/src/expr.rs

@@ -2,10 +2,10 @@ use crate::{
     ast::{
         self, Annotation, ArgBy, ArgName, AssignmentPattern, BinOp, Bls12_381Point,
         ByteArrayFormatPreference, CallArg, Curve, DataType, DataTypeKey, DefinitionLocation,
-        IfBranch, Located, LogicalOpChainKind, ParsedCallArg, Pattern, RecordConstructorArg,
+        Located, LogicalOpChainKind, ParsedCallArg, Pattern, RecordConstructorArg,
         RecordUpdateSpread, Span, TraceKind, TypedArg, TypedAssignmentKind, TypedClause,
-        TypedDataType, TypedRecordUpdateArg, UnOp, UntypedArg, UntypedAssignmentKind,
-        UntypedClause, UntypedRecordUpdateArg,
+        TypedDataType, TypedIfBranch, TypedRecordUpdateArg, UnOp, UntypedArg,
+        UntypedAssignmentKind, UntypedClause, UntypedIfBranch, UntypedRecordUpdateArg,
     },
     builtins::void,
     parser::token::Base,
@@ -127,7 +127,7 @@ pub enum TypedExpr {
     If {
         location: Span,
         #[serde(with = "Vec1Ref")]
-        branches: Vec1<IfBranch<Self>>,
+        branches: Vec1<TypedIfBranch>,
         final_else: Box<Self>,
         tipo: Rc<Type>,
     },
@@ -564,7 +564,7 @@ pub enum UntypedExpr {
 
     If {
         location: Span,
-        branches: Vec1<IfBranch<Self>>,
+        branches: Vec1<UntypedIfBranch>,
         final_else: Box<Self>,
     },
 

crates/aiken-lang/src/format.rs

@@ -2,10 +2,10 @@ use crate::{
     ast::{
         Annotation, ArgBy, ArgName, ArgVia, AssignmentKind, AssignmentPattern, BinOp,
         ByteArrayFormatPreference, CallArg, ClauseGuard, Constant, CurveType, DataType, Definition,
-        Function, IfBranch, LogicalOpChainKind, ModuleConstant, OnTestFailure, Pattern,
-        RecordConstructor, RecordConstructorArg, RecordUpdateSpread, Span, TraceKind, TypeAlias,
-        TypedArg, UnOp, UnqualifiedImport, UntypedArg, UntypedArgVia, UntypedAssignmentKind,
-        UntypedClause, UntypedClauseGuard, UntypedDefinition, UntypedFunction, UntypedModule,
+        Function, LogicalOpChainKind, ModuleConstant, OnTestFailure, Pattern, RecordConstructor,
+        RecordConstructorArg, RecordUpdateSpread, Span, TraceKind, TypeAlias, TypedArg, UnOp,
+        UnqualifiedImport, UntypedArg, UntypedArgVia, UntypedAssignmentKind, UntypedClause,
+        UntypedClauseGuard, UntypedDefinition, UntypedFunction, UntypedIfBranch, UntypedModule,
         UntypedPattern, UntypedRecordUpdateArg, Use, Validator, CAPTURE_VARIABLE,
     },
     docvec,
@@ -703,6 +703,7 @@ impl<'comments> Formatter<'comments> {
         kind: UntypedAssignmentKind,
     ) -> Document<'a> {
         let keyword = match kind {
+            AssignmentKind::Is => unreachable!(),
             AssignmentKind::Let { .. } => "let",
             AssignmentKind::Expect { .. } => "expect",
         };
@@ -1195,7 +1196,7 @@ impl<'comments> Formatter<'comments> {
 
     pub fn if_expr<'a>(
         &mut self,
-        branches: &'a Vec1<IfBranch<UntypedExpr>>,
+        branches: &'a Vec1<UntypedIfBranch>,
         final_else: &'a UntypedExpr,
     ) -> Document<'a> {
         let if_branches = self
@@ -1223,10 +1224,44 @@ impl<'comments> Formatter<'comments> {
     pub fn if_branch<'a>(
         &mut self,
         if_keyword: Document<'a>,
-        branch: &'a IfBranch<UntypedExpr>,
+        branch: &'a UntypedIfBranch,
     ) -> Document<'a> {
         let if_begin = if_keyword
             .append(self.wrap_expr(&branch.condition))
+            .append(match &branch.is {
+                Some(AssignmentPattern {
+                    pattern,
+                    annotation,
+                    ..
+                }) => {
+                    let is_sugar = matches!(
+                        (&pattern, &branch.condition),
+                        (
+                            Pattern::Var { name, .. },
+                            UntypedExpr::Var { name: var_name, .. }
+                        ) if name == var_name
+                    );
+
+                    let Some(annotation) = &annotation else {
+                        unreachable!()
+                    };
+
+                    let is = if is_sugar {
+                        self.annotation(annotation)
+                    } else {
+                        self.pattern(pattern)
+                            .append(": ")
+                            .append(self.annotation(annotation))
+                            .group()
+                    };
+
+                    break_("", " ")
+                        .append("is")
+                        .append(break_("", " "))
+                        .append(is)
+                }
+                None => nil(),
+            })
             .append(break_("{", " {"))
             .group();