egison-5.0.0: hs-src/Language/Egison/Desugar.hs
{-# LANGUAGE TupleSections #-}
{- |
Module : Language.Egison.Desugar
Licence : MIT
This module implements Phase 3-4: Syntactic Desugaring (for untyped path).
For the typed path, desugaring is done inside type inference.
Syntactic Desugaring (Phase 3-4):
- Operator desugaring (infix to function application)
- Anonymous function expansion (cambda: 1#($1 + $2) etc.)
- Match-lambda expansion (convert to match expressions)
- Other syntactic sugar expansions
Design Note (design/implementation.md):
Pattern matching itself is NOT desugared here. Match expressions (IMatchExpr,
IMatchAllExpr) are kept as-is and processed during evaluation (Phase 10).
This allows Egison's sophisticated pattern matching to be implemented in the evaluator.
-}
module Language.Egison.Desugar
( desugarTopExpr
, desugarTopExprs
, desugarExpr
, transVarIndex
) where
import Control.Monad.Except (throwError)
import Data.Char (toUpper)
import Data.Foldable (foldrM)
import Data.List (union)
import Data.Text (pack)
import Language.Egison.AST
import Language.Egison.Data
import Language.Egison.IExpr
import Language.Egison.RState
import Language.Egison.Type.Types (sanitizeMethodName, typeToName, typeConstructorName,
typeExprToType, capitalizeFirst, lowerFirst, TyVar(..))
desugarTopExpr :: TopExpr -> EvalM (Maybe ITopExpr)
desugarTopExpr (Define vwi expr) = do
(var, iexpr) <- desugarDefineWithIndices vwi expr
return . Just $ IDefine var iexpr
desugarTopExpr (DefineWithType typedVwi expr) = do
-- Convert typed definition to regular definition
-- Type information is used for type checking, but the runtime representation is the same
-- Note: Constraints are preserved in the type scheme (by EnvBuilder),
-- and dictionary passing is handled in TypeClassExpand phase
let name = typedVarName typedVwi
indices = typedVarIndices typedVwi
params = typedVarParams typedVwi
vwi = VarWithIndices name indices
-- If there are typed parameters, wrap the body in a lambda
case params of
[] -> do
(var, iexpr) <- desugarDefineWithIndices vwi expr
return . Just $ IDefine var iexpr
_ -> do
-- Create lambda arguments from typed parameters
let argPatterns = map typedParamToArgPattern params
lambdaExpr = LambdaExpr argPatterns expr
(var, iexpr) <- desugarDefineWithIndices vwi lambdaExpr
return . Just $ IDefine var iexpr
desugarTopExpr (Test expr) = Just . ITest <$> desugar expr
desugarTopExpr (Execute expr) = Just . IExecute <$> desugar expr
desugarTopExpr (Load file) = return . Just $ ILoad file
desugarTopExpr (LoadFile file) = return . Just $ ILoadFile file
-- Type class declarations: generate dictionary-passing wrapper functions
-- and register the class methods for dispatch
-- For a class like:
-- class Eq a where
-- (==) (x: a) (y: a) : Bool
-- We generate:
-- 1. Dictionary wrapper: def classEqEq dict x y := (dict_"eq") x y
-- 2. Instance registry variable: def registryEq := {| |}
-- 3. Auto-dispatch function: def autoEqEq x y := (resolveEq x)_"eq" x y
desugarTopExpr (ClassDeclExpr (ClassDecl classNm _typeParams _supers methods)) = do
-- Generate dictionary-passing wrapper functions for each method
methodWrappers <- mapM (desugarClassMethod classNm) methods
-- Generate empty instance registry
let registryDef = makeRegistryDef classNm
case methodWrappers of
[] -> return Nothing
_ -> return $ Just $ IDefineMany (registryDef : methodWrappers)
where
desugarClassMethod :: String -> ClassMethod -> EvalM (Var, IExpr)
desugarClassMethod clsNm (ClassMethod methName methParams _retType _defaultImpl) = do
-- Generate function name: e.g., "classEqEq" for (==) in Eq
let wrapperName = "class" ++ clsNm ++ capitalizeFirst (sanitizeMethodName methName)
var = stringToVar wrapperName
dictVar = "dict"
-- Parameter names: dict, x, y, ...
paramNames = map extractParamName methParams
allParams = dictVar : paramNames
-- Build the body: (dict_"methodName") x y ...
-- dict_"eq" is hash access, then apply to remaining params
let dictAccessExpr = IIndexedExpr False (IVarExpr dictVar)
[Sub (IConstantExpr (StringExpr (pack (sanitizeMethodName methName))))]
bodyExpr = if null paramNames
then dictAccessExpr
else IApplyExpr dictAccessExpr (map IVarExpr paramNames)
lambdaExpr = ILambdaExpr Nothing (map stringToVar allParams) bodyExpr
return (var, lambdaExpr)
-- Create empty instance registry: registryEq := {| |}
makeRegistryDef :: String -> (Var, IExpr)
makeRegistryDef clsNm =
let registryName = "registry" ++ clsNm
var = stringToVar registryName
in (var, IHashExpr [])
extractParamName :: TypedParam -> String
extractParamName (TPVar name _) = name
extractParamName (TPInvertedVar name _) = name
extractParamName (TPUntypedVar name) = name
extractParamName _ = "x" -- fallback
-- Instance declarations: generate a dictionary and individual method definitions
-- For an instance like:
-- instance Eq Integer where
-- (==) x y := x = y
-- (/=) x y := not (x = y)
-- We generate:
-- 1. Individual method functions:
-- def eqIntegerEq x y := x = y
-- def eqIntegerNeq x y := not (x = y)
-- 2. A dictionary for the instance:
-- def eqInteger := {| ("eq", eqIntegerEq), ("neq", eqIntegerNeq) |}
desugarTopExpr (InstanceDeclExpr (InstanceDecl constraints classNm instTypes methods)) = do
-- Check if instTypes is not empty
if null instTypes
then return Nothing
else do
-- Use type constructor name only (without type parameters)
-- e.g., "Collection" not "Collectiona" for [a]
let instTypeName = typeConstructorName (typeExprToType (head instTypes))
-- Generate individual method definitions with constraint parameters
methodDefs <- mapM (desugarInstanceMethod constraints classNm instTypeName) methods
-- Generate dictionary definition (with constraints if any)
let dictDef = makeDictDef constraints classNm instTypeName methods
-- Return all definitions
case methodDefs of
[] -> return Nothing
_ -> return $ Just $ IDefineMany (dictDef : methodDefs)
where
desugarInstanceMethod :: [ConstraintExpr] -> String -> String -> InstanceMethod -> EvalM (Var, IExpr)
desugarInstanceMethod _constrs clsNm typNm (InstanceMethod methName params body) = do
-- Generate function name using type constructor name only
-- e.g., "eqCollectionEq" not "eqCollectionaEq" for instance {Eq a} Eq [a]
let funcName = lowerFirst clsNm ++ typNm ++ capitalizeFirst (sanitizeMethodName methName)
var = stringToVar funcName
-- Do NOT add dictionary parameters here!
-- Dictionary parameters will be added automatically by addDictionaryParametersT
-- after type inference, based on the inferred constraints.
-- This allows the method body to be properly type-checked with constraints.
-- Create lambda expression with only the method parameters
let lambdaArgs = map (\p -> Arg (APPatVar (VarWithIndices p []))) params
lambdaExpr = if null params then body else LambdaExpr lambdaArgs body
iexpr <- desugar lambdaExpr
return (var, iexpr)
makeDictDef :: [ConstraintExpr] -> String -> String -> [InstanceMethod] -> (Var, IExpr)
makeDictDef _constrs clsNm typNm meths =
let dictName = lowerFirst clsNm ++ typNm -- e.g., "eqCollection"
dictVar = stringToVar dictName
-- For nested instances (with constraints), the dictionary becomes a function
-- that takes dictionary parameters and returns a hash.
-- e.g., for instance {Eq a} Eq [a]:
-- eqCollection = \dict_Eq -> {| ("eq", eqCollectionEq dict_Eq), ... |}
--
-- Dictionary parameters will be automatically added by addDictionaryParametersT
-- after type inference, so we don't add them here manually.
-- We just create the hash with references to the methods.
hashEntries = map (makeHashEntry clsNm typNm) meths
hashExpr = IHashExpr hashEntries
in (dictVar, hashExpr)
makeHashEntry :: String -> String -> InstanceMethod -> (IExpr, IExpr)
makeHashEntry clsNm typNm (InstanceMethod methName _ _) =
let keyExpr = IConstantExpr (StringExpr (pack (sanitizeMethodName methName)))
-- Reference to the method function
funcName = lowerFirst clsNm ++ typNm ++ capitalizeFirst (sanitizeMethodName methName)
valueExpr = IVarExpr funcName
in (keyExpr, valueExpr)
-- Inductive declarations don't produce runtime code
-- Constructor registration is handled by the type system
desugarTopExpr (InductiveDecl _ _ _) = return Nothing
-- Infix declarations don't produce runtime code
desugarTopExpr (InfixDecl _ _) = return Nothing
desugarTopExpr (PatternInductiveDecl _ _ _) = return Nothing -- Handled in environment building phase
-- Pattern function declarations need type checking, so convert to IPatternFunctionDecl
desugarTopExpr (PatternFunctionDecl name typeParams params retType body) = do
let paramTypes = map (\(pname, pty) -> (pname, typeExprToType pty)) params
retType' = typeExprToType retType
tyVars = map TyVar typeParams
body' <- desugarPattern body
return . Just $ IPatternFunctionDecl name tyVars paramTypes retType' body'
-- Symbol declarations
desugarTopExpr (DeclareSymbol names mTypeExpr) = do
-- Convert type expression to type (defaults to Integer if not specified)
let ty = case mTypeExpr of
Just texpr -> typeExprToType texpr
Nothing -> typeExprToType TEInt
return . Just $ IDeclareSymbol names (Just ty)
-- | Convert TypedParam to Arg ArgPattern for lambda expressions
typedParamToArgPattern :: TypedParam -> Arg ArgPattern
typedParamToArgPattern (TPVar pname _) =
Arg (APPatVar (VarWithIndices pname []))
typedParamToArgPattern (TPInvertedVar pname _) =
InvertedArg (APPatVar (VarWithIndices pname []))
typedParamToArgPattern (TPTuple elems) =
Arg (APTuplePat (map typedParamToArgPattern elems))
typedParamToArgPattern (TPWildcard _) =
Arg APWildCard
typedParamToArgPattern (TPUntypedVar pname) =
Arg (APPatVar (VarWithIndices pname []))
typedParamToArgPattern TPUntypedWildcard =
Arg APWildCard
desugarTopExprs :: [TopExpr] -> EvalM [ITopExpr]
desugarTopExprs [] = return []
desugarTopExprs (expr : exprs) = do
expr' <- desugarTopExpr expr
case expr' of
Nothing -> desugarTopExprs exprs
Just expr' -> (expr' :) <$> desugarTopExprs exprs
desugarExpr :: Expr -> EvalM IExpr
desugarExpr = desugar
desugar :: Expr -> EvalM IExpr
desugar (ConstantExpr c) = return $ IConstantExpr c
desugar (VarExpr var) = return $ IVarExpr var
desugar (AlgebraicDataMatcherExpr patterns) = do
matcherName <- fresh
let matcherRef = IVarExpr matcherName
matcher <- genMatcherClauses patterns matcherRef
return $ ILetRecExpr [(PDPatVar (stringToVar matcherName), matcher)] matcherRef
where
genMatcherClauses :: [(String, [Expr])] -> IExpr -> EvalM IExpr
genMatcherClauses patterns matcher = do
main <- genMainClause patterns matcher
body <- mapM genMatcherClause patterns
footer <- genSomethingClause
let clauses = [main] ++ body ++ [footer]
return $ IMatcherExpr clauses
genMainClause :: [(String, [Expr])] -> IExpr -> EvalM (PrimitivePatPattern, IExpr, [(IPrimitiveDataPattern, IExpr)])
genMainClause patterns matcher = do
clauses <- genClauses patterns
return (PPValuePat "val", ITupleExpr [],
[(PDPatVar (stringToVar "tgt"),
IMatchExpr BFSMode
(ITupleExpr [IVarExpr "val", IVarExpr "tgt"])
(ITupleExpr [matcher, matcher])
clauses)])
where
genClauses :: [(String, [Expr])] -> EvalM [IMatchClause]
genClauses patterns = (++) <$> mapM genClause patterns
<*> pure [(ITuplePat [IWildCard, IWildCard], matchingFailure)]
genClause :: (String, [Expr]) -> EvalM IMatchClause
genClause pattern = do
(pat0, pat1) <- genMatchingPattern pattern
return (ITuplePat [pat0, pat1], matchingSuccess)
genMatchingPattern :: (String, [Expr]) -> EvalM (IPattern, IPattern)
genMatchingPattern (name, patterns) = do
names <- mapM (const fresh) patterns
return (IInductivePat name (map IPatVar names),
IInductivePat name (map (IValuePat . IVarExpr) names))
genMatcherClause :: (String, [Expr]) -> EvalM (PrimitivePatPattern, IExpr, [(IPrimitiveDataPattern, IExpr)])
genMatcherClause pattern = do
(ppat, matchers) <- genPrimitivePatPat pattern
(dpat, body) <- genPrimitiveDataPat pattern
return (ppat, ITupleExpr matchers, [(dpat, ICollectionExpr [ITupleExpr body]), (PDWildCard, matchingFailure)])
where
genPrimitivePatPat :: (String, [Expr]) -> EvalM (PrimitivePatPattern, [IExpr])
genPrimitivePatPat (name, matchers) = do
patterns' <- mapM (const $ return PPPatVar) matchers
matchers' <- mapM desugar matchers
return (PPInductivePat name patterns', matchers')
genPrimitiveDataPat :: (String, [Expr]) -> EvalM (IPrimitiveDataPattern, [IExpr])
genPrimitiveDataPat (name, patterns) = do
patterns' <- mapM (const fresh) patterns
return (PDInductivePat (capitalize name) $ map (PDPatVar . stringToVar) patterns', map IVarExpr patterns')
capitalize :: String -> String
capitalize (x:xs) = toUpper x : xs
genSomethingClause :: EvalM (PrimitivePatPattern, IExpr, [(IPrimitiveDataPattern, IExpr)])
genSomethingClause =
return (PPPatVar, ITupleExpr [IConstantExpr SomethingExpr], [(PDPatVar (stringToVar "tgt"), ICollectionExpr [IVarExpr "tgt"])])
matchingSuccess :: IExpr
matchingSuccess = ICollectionExpr [ITupleExpr []]
matchingFailure :: IExpr
matchingFailure = ICollectionExpr []
desugar (MatchAllLambdaExpr matcher clauses) = do
name <- fresh
ILambdaExpr Nothing [stringToVar name] <$>
desugar (MatchAllExpr BFSMode (VarExpr name) matcher clauses)
desugar (MatchLambdaExpr matcher clauses) = do
name <- fresh
ILambdaExpr Nothing [stringToVar name] <$>
desugar (MatchExpr BFSMode (VarExpr name) matcher clauses)
desugar (IndexedExpr override expr indices) = do
expr' <- desugar expr
desugarIndexedExpr override expr' indices
where
desugarIndexedExpr :: Bool -> IExpr -> [IndexExpr Expr] -> EvalM IExpr
desugarIndexedExpr override expr' indices =
case indices of
[] -> return expr'
(MultiSubscript x y:indices') ->
case (x, y) of
(IndexedExpr override1 e1 [n1], IndexedExpr _ _ [n2]) -> do
expr'' <- desugarMultiScript override expr' ISubrefsExpr override1 e1 n1 n2
desugarIndexedExpr False expr'' indices'
_ -> throwError $ Default "Index should be IndexedExpr for multi subscript"
(MultiSuperscript x y:indices') ->
case (x, y) of
(IndexedExpr override1 e1 [n1], IndexedExpr _ _ [n2]) -> do
expr'' <- desugarMultiScript override expr' ISuprefsExpr override1 e1 n1 n2
desugarIndexedExpr False expr'' indices'
_ -> throwError $ Default "Index should be IndexedExpr for multi superscript"
_ -> do
let (is, indices') = break isMulti indices
expr'' <- IIndexedExpr override expr' <$> mapM desugarIndex is
desugarIndexedExpr False expr'' indices'
desugarMultiScript override expr' refExpr override1 e1 n1 n2 = do
k <- fresh
n1' <- desugar (extractIndexExpr n1)
n2' <- desugar (extractIndexExpr n2)
e1' <- desugar e1
return $ refExpr override expr' (makeIApply "map"
[ILambdaExpr Nothing [stringToVar k] (IIndexedExpr override1 e1' [Sub (IVarExpr k)]),
makeIApply "between" [n1', n2']])
isMulti (MultiSubscript _ _) = True
isMulti (MultiSuperscript _ _) = True
isMulti _ = False
desugar (SubrefsExpr bool expr1 expr2) =
ISubrefsExpr bool <$> desugar expr1 <*> desugar expr2
desugar (SuprefsExpr bool expr1 expr2) =
ISuprefsExpr bool <$> desugar expr1 <*> desugar expr2
desugar (UserrefsExpr bool expr1 expr2) =
IUserrefsExpr bool <$> desugar expr1 <*> desugar expr2
desugar (TupleExpr exprs) = ITupleExpr <$> mapM desugar exprs
desugar (CollectionExpr xs) = ICollectionExpr <$> mapM desugar xs
desugar (ConsExpr x xs) = IConsExpr <$> desugar x <*> desugar xs
desugar (JoinExpr x xs) = IJoinExpr <$> desugar x <*> desugar xs
desugar (HashExpr exprPairs) =
IHashExpr <$> mapM (\(expr1, expr2) -> (,) <$> desugar expr1 <*> desugar expr2) exprPairs
desugar (VectorExpr exprs) =
IVectorExpr <$> mapM desugar exprs
desugar (TensorExpr nsExpr xsExpr) =
ITensorExpr <$> desugar nsExpr <*> desugar xsExpr
-- Desugar of LambdaExpr takes place in 2 stages.
-- * LambdaExpr -> LambdaExpr' : Desugar pattern matches at the arg positions
-- * LambdaExpr' -> ILambdaExpr : Desugar Arg and InvertedArg
desugar (LambdaExpr args expr) = do
(args', expr') <- foldrM desugarArg ([], expr) args
desugar $ LambdaExpr' args' expr'
where
desugarArg :: Arg ArgPattern -> ([Arg VarWithIndices], Expr) -> EvalM ([Arg VarWithIndices], Expr)
desugarArg (Arg x) (args, expr) = do
(var, expr') <- desugarArgPat x expr
return (Arg var : args, expr')
desugarArg (InvertedArg x) (args, expr) = do
(var, expr') <- desugarArgPat x expr
return (InvertedArg var : args, expr')
-- Desugar argument patterns. Examples:
-- \$(%x, %y) -> expr ==> \$tmp -> let (tmp1, tmp2) := tmp in (\%x %y -> expr) tmp1 tmp2
-- \(x, (y, z)) -> expr ==> \tmp -> let (tmp1, tmp2) := tmp in (\x (y, z) -> expr) tmp1 tmp2
-- \%($x :: xs) -> expr ==> \%tmp -> let (tmp1 :: xs) := tmp in (\$x %xs -> expr) tmp1 tmp2
desugarArgPat :: ArgPattern -> Expr -> EvalM (VarWithIndices, Expr)
desugarArgPat APWildCard expr = do
tmp <- fresh
let tmp' = stringToVarWithIndices tmp
return (tmp', LetExpr [Bind PDWildCard (VarExpr tmp)] expr)
desugarArgPat (APPatVar var) expr = return (var, expr)
desugarArgPat (APTuplePat args) expr = do
tmp <- fresh
let tmp' = stringToVarWithIndices tmp
tmps <- mapM (const fresh) args
return (tmp', LetExpr [Bind (PDTuplePat (map PDPatVar tmps)) (VarExpr tmp)]
(ApplyExpr (LambdaExpr args expr) (map VarExpr tmps)))
desugarArgPat (APInductivePat ctor args) expr = do
tmp <- fresh
let tmp' = stringToVarWithIndices tmp
tmps <- mapM (const fresh) args
return (tmp', LetExpr [Bind (PDInductivePat ctor (map PDPatVar tmps)) (VarExpr tmp)]
(ApplyExpr (LambdaExpr args expr) (map VarExpr tmps)))
desugarArgPat APEmptyPat expr = do
tmp <- fresh
let tmp' = stringToVarWithIndices tmp
return (tmp', LetExpr [Bind PDEmptyPat (VarExpr tmp)] expr)
desugarArgPat (APConsPat arg1 arg2) expr = do
tmp <- fresh
let tmp' = stringToVarWithIndices tmp
tmp1 <- fresh
tmp2 <- fresh
return (tmp', LetExpr [Bind (PDConsPat (PDPatVar tmp1) (PDPatVar tmp2)) (VarExpr tmp)]
(ApplyExpr (LambdaExpr [arg1, Arg arg2] expr) [VarExpr tmp1, VarExpr tmp2]))
desugarArgPat (APSnocPat arg1 arg2) expr = do
tmp <- fresh
let tmp' = stringToVarWithIndices tmp
tmp1 <- fresh
tmp2 <- fresh
return (tmp', LetExpr [Bind (PDSnocPat (PDPatVar tmp1) (PDPatVar tmp2)) (VarExpr tmp)]
(ApplyExpr (LambdaExpr [Arg arg1, arg2] expr) [VarExpr tmp1, VarExpr tmp2]))
desugar (LambdaExpr' vwis expr) = do
let (vwis', expr') = foldr desugarInvertedArgs ([], expr) vwis
let args' = map varWithIndicesToVar vwis'
expr' <- desugar expr'
return $ ILambdaExpr Nothing args' expr'
where
desugarInvertedArgs :: Arg VarWithIndices -> ([VarWithIndices], Expr) -> ([VarWithIndices], Expr)
desugarInvertedArgs (Arg x) (args, expr) = (x : args, expr)
desugarInvertedArgs (InvertedArg x) (args, expr) =
let varName = extractNameFromVarWithIndices x
flippedExpr = FlipIndicesExpr (VarExpr varName)
bindPat = PDPatVar varName
in (x : args, LetExpr [Bind bindPat flippedExpr] expr)
desugar (MemoizedLambdaExpr names expr) =
IMemoizedLambdaExpr names <$> desugar expr
-- Typed memoized lambda is desugared the same way (type info used only for type checking)
desugar (TypedMemoizedLambdaExpr params _ body) =
IMemoizedLambdaExpr (extractParamNames params) <$> desugar body
where
extractParamNames = concatMap extractName
extractName (TPVar name _) = [name]
extractName (TPInvertedVar name _) = [name]
extractName (TPTuple elems) = concatMap extractName elems
extractName (TPWildcard _) = []
extractName (TPUntypedVar name) = [name]
extractName TPUntypedWildcard = []
desugar (CambdaExpr name expr) =
ICambdaExpr name <$> desugar expr
desugar (PatternFunctionExpr _names _pattern) =
-- Pattern functions are only defined at TopExpr level
-- They should not appear in expression context
throwError $ Default "Pattern functions cannot be used as expressions"
desugar (IfExpr expr0 expr1 expr2) =
IIfExpr <$> desugar expr0 <*> desugar expr1 <*> desugar expr2
desugar (LetExpr binds expr) =
ILetExpr <$> desugarBindings binds <*> desugar expr
desugar (LetRecExpr binds expr) =
ILetRecExpr <$> desugarBindings binds <*> desugar expr
desugar (WithSymbolsExpr vars expr) =
IWithSymbolsExpr vars <$> desugar expr
desugar (MatchExpr pmmode expr0 expr1 clauses) =
IMatchExpr pmmode <$> desugar expr0 <*> desugar expr1 <*> desugarMatchClauses clauses
desugar (MatchAllExpr pmmode expr0 expr1 clauses) =
IMatchAllExpr pmmode <$> desugar expr0 <*> desugar expr1 <*> desugarMatchClauses clauses
desugar (DoExpr binds expr) =
IDoExpr <$> desugarBindings binds <*> desugar expr
desugar (PrefixExpr "-" expr) = do
expr' <- desugar expr
return $ makeIApply "*" [IConstantExpr (IntegerExpr (-1)), expr']
desugar (PrefixExpr "!" (ApplyExpr expr args)) =
IWedgeApplyExpr <$> desugar expr <*> mapM desugar args
desugar (PrefixExpr "'" expr) = IQuoteExpr <$> desugar expr
desugar (PrefixExpr "`" expr) = IQuoteSymbolExpr <$> desugar expr
desugar (PrefixExpr op _) = fail ("Unknown prefix " ++ op)
desugar (InfixExpr op expr1 expr2) | isWedge op =
(\x y -> IWedgeApplyExpr (IVarExpr (repr op)) [x, y])
<$> desugar expr1 <*> desugar expr2
desugar (InfixExpr op expr1 expr2) | repr op == "::" =
IConsExpr <$> desugar expr1 <*> desugar expr2
desugar (InfixExpr op expr1 expr2) | repr op == "++" =
IJoinExpr <$> desugar expr1 <*> desugar expr2
desugar (InfixExpr op expr1 expr2) =
(\x y -> makeIApply (repr op) [x, y]) <$> desugar expr1 <*> desugar expr2
-- section
--
-- If `op` is not a cambda, simply desugar it into the function
desugar (SectionExpr op Nothing Nothing)
| not (isWedge op || repr op `elem` ["::", "++"]) =
desugar (VarExpr (repr op))
desugar (SectionExpr op Nothing Nothing) = do
x <- fresh
y <- fresh
ILambdaExpr Nothing [stringToVar x, stringToVar y] <$> desugar (InfixExpr op (VarExpr x) (VarExpr y))
desugar (SectionExpr op Nothing (Just expr2)) = do
x <- fresh
ILambdaExpr Nothing [stringToVar x] <$> desugar (InfixExpr op (VarExpr x) expr2)
desugar (SectionExpr op (Just expr1) Nothing) = do
y <- fresh
ILambdaExpr Nothing [stringToVar y] <$> desugar (InfixExpr op expr1 (VarExpr y))
desugar SectionExpr{} = throwError $ Default "Cannot reach here: section with both arguments"
desugar (SeqExpr expr0 expr1) =
ISeqExpr <$> desugar expr0 <*> desugar expr1
desugar (GenerateTensorExpr fnExpr sizeExpr) =
IGenerateTensorExpr <$> desugar fnExpr <*> desugar sizeExpr
desugar (TensorContractExpr tExpr) =
ITensorContractExpr <$> desugar tExpr
desugar (TensorMapExpr (LambdaExpr' [x] (TensorMapExpr (LambdaExpr' [y] expr) b)) a) =
desugar (TensorMap2Expr (LambdaExpr' [x, y] expr) a b)
desugar (TensorMapExpr (LambdaExpr [x] (TensorMapExpr (LambdaExpr [y] expr) b)) a) =
desugar (TensorMap2Expr (LambdaExpr [x, y] expr) a b)
desugar (TensorMapExpr fnExpr tExpr) =
ITensorMapExpr <$> desugar fnExpr <*> desugar tExpr
desugar (TensorMap2Expr fnExpr t1Expr t2Expr) =
ITensorMap2Expr <$> desugar fnExpr <*> desugar t1Expr <*> desugar t2Expr
desugar (TransposeExpr vars expr) =
-- ITransposeExpr takes (permutation, tensor) as arguments to match tTranspose
ITransposeExpr <$> desugar vars <*> desugar expr
desugar (FlipIndicesExpr expr) =
IFlipIndicesExpr <$> desugar expr
desugar (ApplyExpr expr args) =
IApplyExpr <$> desugar expr <*> mapM desugar args
desugar FreshVarExpr = do
id <- fresh
return $ IVarExpr (":::" ++ id)
desugar (MatcherExpr patternDefs) =
IMatcherExpr <$> mapM desugarPatternDef patternDefs
desugar (AnonParamExpr n) = return $ IVarExpr ('%' : show n)
desugar (AnonParamFuncExpr n expr) = do
let args = map (\n -> stringToVarWithIndices ('%' : show n)) [1..n]
lambda <- desugar $ LambdaExpr' (map Arg args) expr
return $ ILetRecExpr [(PDPatVar (stringToVar "%0"), lambda)] (IVarExpr "%0")
desugar (AnonTupleParamFuncExpr 1 expr) = do
lambda <- desugar $ LambdaExpr' [Arg (stringToVarWithIndices "%1")] expr
return $ ILetRecExpr [(PDPatVar (stringToVar "%0"), lambda)] (IVarExpr "%0")
desugar (AnonTupleParamFuncExpr n expr) = do
let args = map (\n -> stringToVarWithIndices ('%' : show n)) [1..n]
lambda <- desugar $
LambdaExpr [Arg (APTuplePat $ map (Arg . APPatVar) args)] expr
return $ ILetRecExpr [(PDPatVar (stringToVar "%0"), lambda)] (IVarExpr "%0")
desugar (AnonListParamFuncExpr n expr) = do
let args' = map (\n -> Arg (APPatVar (stringToVarWithIndices ('%' : show n)))) [1..n]
let args = foldr APConsPat APEmptyPat args'
lambda <- desugar $ LambdaExpr [Arg args] expr
return $ ILetRecExpr [(PDPatVar (stringToVar "%0"), lambda)] (IVarExpr "%0")
desugar (QuoteExpr expr) =
IQuoteExpr <$> desugar expr
desugar (QuoteSymbolExpr expr) =
IQuoteSymbolExpr <$> desugar expr
desugar (WedgeApplyExpr expr args) =
IWedgeApplyExpr <$> desugar expr <*> mapM desugar args
desugar (FunctionExpr args) = return $ IFunctionExpr args
-- Type annotation is erased at runtime
desugar (TypeAnnotation expr _typeExpr) = desugar expr
-- Typed lambda is desugared to regular lambda
desugar (TypedLambdaExpr params _retType body) = do
let args = map (\(name, _) -> Arg (APPatVar (VarWithIndices name []))) params
desugar $ LambdaExpr args body
desugarIndex :: IndexExpr Expr -> EvalM (Index IExpr)
desugarIndex (Subscript e) = Sub <$> desugar e
desugarIndex (Superscript e) = Sup <$> desugar e
desugarIndex (SupSubscript e) = SupSub <$> desugar e
desugarIndex (Userscript e) = User <$> desugar e
desugarIndex _ = undefined
desugarPattern :: Pattern -> EvalM IPattern
desugarPattern pat =
case collectName pat of
[] -> desugarPattern' pat
names -> ILetPat (map makeBinding names) <$> desugarPattern' pat
where
collectNames :: [Pattern] -> [String]
collectNames pats = foldl union [] (map collectName pats)
collectName :: Pattern -> [String]
collectName (ForallPat pat1 pat2) = collectName pat1 `union` collectName pat2
collectName (InfixPat _ pat1 pat2) = collectName pat1 `union` collectName pat2
collectName (NotPat pat) = collectName pat
collectName (AndPat pat1 pat2) = collectName pat1 `union` collectName pat2
collectName (OrPat pat1 pat2) = collectName pat1 `union` collectName pat2
collectName (TuplePat pats) = collectNames pats
collectName (InductiveOrPApplyPat _ pats) = collectNames pats
collectName (InductivePat _ pats) = collectNames pats
collectName (PApplyPat _ pats) = collectNames pats
collectName (DApplyPat _ pats) = collectNames pats
collectName (LoopPat _ (LoopRange _ _ endNumPat) pat1 pat2) = collectName endNumPat `union` collectName pat1 `union` collectName pat2
collectName (LetPat _ pat) = collectName pat
collectName (IndexedPat (PatVar var) _) = [var]
collectName _ = []
makeBinding :: String -> IBindingExpr
makeBinding var = (PDPatVar (stringToVar var), IHashExpr [])
desugarPattern' :: Pattern -> EvalM IPattern
desugarPattern' WildCard = return IWildCard
desugarPattern' ContPat = return IContPat
desugarPattern' SeqNilPat = return ISeqNilPat
desugarPattern' LaterPatVar = return ILaterPatVar
desugarPattern' (VarPat v) = return (IVarPat v)
desugarPattern' (PatVar var) = return (IPatVar var)
desugarPattern' (ValuePat expr) = IValuePat <$> desugar expr
desugarPattern' (PredPat expr) = IPredPat <$> desugar expr
desugarPattern' (NotPat pat) = INotPat <$> desugarPattern' pat
desugarPattern' (AndPat pat1 pat2) = IAndPat <$> desugarPattern' pat1 <*> desugarPattern' pat2
desugarPattern' (OrPat pat1 pat2) = IOrPat <$> desugarPattern' pat1 <*> desugarPattern' pat2
desugarPattern' (ForallPat pat1 pat2) = IForallPat <$> desugarPattern' pat1 <*> desugarPattern' pat2
desugarPattern' (InfixPat Op{ repr = "&" } pat1 pat2) =
IAndPat <$> desugarPattern' pat1 <*> desugarPattern' pat2
desugarPattern' (InfixPat Op{ repr = "|" } pat1 pat2) =
IOrPat <$> desugarPattern' pat1 <*> desugarPattern' pat2
desugarPattern' (InfixPat Op{ repr = f } pat1 pat2) =
(\x y -> IInductivePat f [x, y]) <$> desugarPattern' pat1 <*> desugarPattern' pat2
desugarPattern' (TuplePat pats) = ITuplePat <$> mapM desugarPattern' pats
desugarPattern' (InductiveOrPApplyPat name pats) = IInductiveOrPApplyPat name <$> mapM desugarPattern' pats
-- Convert all InductivePat to IInductiveOrPApplyPat since we cannot distinguish between
-- pattern constructors and pattern functions at parse time
desugarPattern' (InductivePat name pats) = IInductiveOrPApplyPat name <$> mapM desugarPattern' pats
desugarPattern' (IndexedPat pat exprs) = IIndexedPat <$> desugarPattern' pat <*> mapM desugar exprs
desugarPattern' (PApplyPat expr pats) = IPApplyPat <$> desugar expr <*> mapM desugarPattern' pats
desugarPattern' (DApplyPat pat pats) = IDApplyPat <$> desugarPattern' pat <*> mapM desugarPattern' pats
desugarPattern' (LoopPat name range pat1 pat2) = ILoopPat name <$> desugarLoopRange range <*> desugarPattern' pat1 <*> desugarPattern' pat2
desugarPattern' (LetPat binds pat) = ILetPat <$> desugarBindings binds <*> desugarPattern' pat
desugarPattern' (SeqConsPat pat1 pat2) = ISeqConsPat <$> desugarPattern' pat1 <*> desugarPattern' pat2
desugarLoopRange :: LoopRange -> EvalM ILoopRange
desugarLoopRange (LoopRange sExpr eExpr pat) =
ILoopRange <$> desugar sExpr <*> desugar eExpr <*> desugarPattern' pat
desugarBindings :: [BindingExpr] -> EvalM [IBindingExpr]
desugarBindings = mapM desugarBinding
where
desugarBinding (Bind name expr) = do
let name' = fmap stringToVar name
expr' <- desugar expr
case (name, expr') of
(PDPatVar var, ILambdaExpr Nothing args body) ->
return (name', ILambdaExpr (Just (Var var [])) args body)
_ -> return (name', expr')
desugarBinding (BindWithIndices vwi expr) = do
(var, iexpr) <- desugarDefineWithIndices vwi expr
return (PDPatVar var, iexpr)
-- BindWithType: desugar like DefineWithType
desugarBinding (BindWithType typedVarWI body) = do
let name = typedVarName typedVarWI
params = typedVarParams typedVarWI
argPatterns = map typedParamToArgPattern params
lambdaExpr = if null argPatterns
then body
else LambdaExpr argPatterns body
body' <- desugar lambdaExpr
let body'' = case body' of
ILambdaExpr Nothing args b -> ILambdaExpr (Just (Var name [])) args b
other -> other
return (PDPatVar (Var name []), body'')
desugarMatchClauses :: [MatchClause] -> EvalM [IMatchClause]
desugarMatchClauses = mapM (\(pat, expr) -> (,) <$> desugarPattern pat <*> desugar expr)
desugarPatternDef :: PatternDef -> EvalM IPatternDef
desugarPatternDef (PatternDef pp matcher pds) =
(pp,,) <$> desugar matcher <*> desugarPrimitiveDataMatchClauses pds
desugarPrimitiveDataMatchClauses :: [(PrimitiveDataPattern, Expr)] -> EvalM [(IPrimitiveDataPattern, IExpr)]
desugarPrimitiveDataMatchClauses = mapM (\(pd, expr) -> (fmap stringToVar pd,) <$> desugar expr)
desugarDefineWithIndices :: VarWithIndices -> Expr -> EvalM (Var, IExpr)
-- Case 1: No indices - simple desugaring without withSymbols/transpose
desugarDefineWithIndices (VarWithIndices name []) expr = do
expr' <- desugar expr
return (Var name [], expr')
-- Case 2: Non-empty indices - wrap with withSymbols and transpose
desugarDefineWithIndices (VarWithIndices name is) expr = do
let (isSubs, indexNames) = unzip $ concatMap extractSubSupIndex is
expr <- if any isExtendedIndice is
then desugarExtendedIndices is isSubs indexNames expr
else return expr
body <- desugar expr
let indexNamesCollection = ICollectionExpr (map IVarExpr indexNames)
let is' = map (\b -> if b then Sub Nothing else Sup Nothing) isSubs
-- ITransposeExpr takes (permutation, tensor) as arguments to match tTranspose
return (Var name is', IWithSymbolsExpr indexNames (ITransposeExpr indexNamesCollection body))
varWithIndicesToVar :: VarWithIndices -> Var
varWithIndicesToVar (VarWithIndices name is) = Var name (concatMap transVarIndex is)
transVarIndex :: VarIndex -> [Index (Maybe Var)]
transVarIndex (VSubscript x) = [Sub (Just (stringToVar x))]
transVarIndex (VSuperscript x) = [Sup (Just (stringToVar x))]
transVarIndex (VMultiSubscript x s e) = [MultiSub (Just (stringToVar x)) s (Just (stringToVar e))]
transVarIndex (VMultiSuperscript x s e) = [MultiSup (Just (stringToVar x)) s (Just (stringToVar e))]
transVarIndex (VGroupScripts xs) = concatMap transVarIndex xs
transVarIndex (VSymmScripts xs) = concatMap transVarIndex xs
transVarIndex (VAntiSymmScripts xs) = concatMap transVarIndex xs
extractSubSupIndex :: VarIndex -> [(Bool, String)]
extractSubSupIndex (VSubscript x) = [(True, x)]
extractSubSupIndex (VSuperscript x) = [(False, x)]
extractSubSupIndex (VGroupScripts xs) = concatMap extractSubSupIndex xs
extractSubSupIndex (VSymmScripts xs) = concatMap extractSubSupIndex xs
extractSubSupIndex (VAntiSymmScripts xs) = concatMap extractSubSupIndex xs
desugarExtendedIndices :: [VarIndex] -> [Bool] -> [String] -> Expr -> EvalM Expr
desugarExtendedIndices indices isSubs indexNames tensorBody = do
tensorName <- fresh
tensorGenExpr <- f indices (VarExpr tensorName) [] []
let indexFunctionExpr = LambdaExpr [Arg $ foldr APConsPat APEmptyPat (map (Arg . APPatVar) (map stringToVarWithIndices indexNames))] tensorGenExpr
let genTensorExpr = GenerateTensorExpr indexFunctionExpr (makeApply "tensorShape" [VarExpr tensorName])
let tensorIndices = zipWith (\isSub name -> if isSub then Subscript (VarExpr name) else Superscript (VarExpr name)) isSubs indexNames
return $ LetExpr [Bind (PDPatVar tensorName) tensorBody] (IndexedExpr True genTensorExpr tensorIndices)
where
f :: [VarIndex] -> Expr -> [String] -> [BindingExpr] -> EvalM Expr
f [] expr [] [] = return expr
f [] expr [] bindings = return $ LetRecExpr bindings expr
f [] expr signs bindings =
return $ LetRecExpr bindings (makeApply "product" [CollectionExpr (map VarExpr signs ++ [expr])])
f (index:indices) expr signs bindings = do
(indices', signs', bindings') <- genBindings index
let isSubs = subOrSupScripts index
symbols <- mapM (const fresh) isSubs
let is = zipWith (\x isSub -> (if isSub then Subscript else Superscript) (VarExpr x)) symbols isSubs
f indices (IndexedExpr True expr is)
(signs ++ signs') (bindings ++ bindings' ++ [Bind (foldr (PDConsPat . PDPatVar) PDEmptyPat symbols) indices'])
subOrSupScripts :: VarIndex -> [Bool]
subOrSupScripts VSubscript{} = [True]
subOrSupScripts VSuperscript{} = [False]
subOrSupScripts (VGroupScripts xs) = concatMap subOrSupScripts xs
subOrSupScripts (VSymmScripts xs) = concatMap subOrSupScripts xs
subOrSupScripts (VAntiSymmScripts xs) = concatMap subOrSupScripts xs
genBindings :: VarIndex -> EvalM (Expr, [String], [BindingExpr])
genBindings (VSubscript x) = return (CollectionExpr [VarExpr x], [], [])
genBindings (VSuperscript x) = return (CollectionExpr [VarExpr x], [], [])
genBindings (VGroupScripts xs) = do
(indices, signss, bindingss) <- unzip3 <$> mapM genBindings xs
let newIndices =
-- If indices are all CollectionExpr, we can calculate the concatenated result of them
case allCollections indices of
Just xs -> CollectionExpr xs
Nothing -> makeApply "concat" [CollectionExpr indices]
return (newIndices, concat signss, concat bindingss)
where
allCollections [] = Just []
allCollections (CollectionExpr xs : exprs) = (xs ++) <$> allCollections exprs
allCollections _ = Nothing
genBindings (VSymmScripts xs) = do
(indices, signss, bindingss) <- unzip3 <$> mapM genBindings xs
let signs = concat signss
let bindings = concat bindingss
sortedCollectionName <- fresh
let newBindings = bindings ++ [Bind (PDTuplePat [PDWildCard, PDPatVar sortedCollectionName]) (makeApply "sortWithSign" [CollectionExpr indices])]
return (VarExpr sortedCollectionName, signs, newBindings)
genBindings (VAntiSymmScripts xs) = do
(indices, signss, bindingss) <- unzip3 <$> mapM genBindings xs
let signs = concat signss
let bindings = concat bindingss
sortedCollectionName <- fresh
signName <- fresh
let newBindings = bindings ++ [Bind (PDTuplePat [PDPatVar signName, PDPatVar sortedCollectionName]) (makeApply "sortWithSign" [CollectionExpr indices])]
return (VarExpr sortedCollectionName, signName : signs, newBindings)
--
-- Utils
--
extractIndexExpr :: IndexExpr a -> a
extractIndexExpr (Subscript x) = x
extractIndexExpr (Superscript x) = x
extractIndexExpr (SupSubscript x) = x
extractIndexExpr (Userscript x) = x
extractIndexExpr _ = error "extractIndexExpr: Not supported"
isExtendedIndice :: VarIndex -> Bool
isExtendedIndice VSubscript{} = False
isExtendedIndice VSuperscript{} = False
isExtendedIndice (VGroupScripts xs) = isExtendedIndice (head xs)
isExtendedIndice _ = True