+ -- Checks whether an expression contain a non-scalar subexpression.
+ --
+ -- Precodition: The variables in the first argument are scalar.
+ --
+ -- In case of a recursive binding group, we /assume/ that all bindings are scalar (by adding
+ -- them to the list of scalar variables) and then check them. If one of them turns out not to
+ -- be scalar, the entire group is regarded as not being scalar.
+ --
+ -- FIXME: Currently, doesn't regard external (non-data constructor) variable and anonymous
+ -- data constructor as scalar. Should be changed once scalar types are passed
+ -- through VectInfo.
+ --
+ is_scalar :: VarSet -> CoreExpr -> Bool
+ is_scalar scalars (Var v) = v `elemVarSet` scalars
+ is_scalar _scalars (Lit _) = True
+ is_scalar scalars e@(App e1 e2)
+ | maybe_parr_ty (exprType e) = False
+ | otherwise = is_scalar scalars e1 && is_scalar scalars e2
+ is_scalar scalars (Lam var body)
+ | maybe_parr_ty (varType var) = False
+ | otherwise = is_scalar (scalars `extendVarSet` var) body
+ is_scalar scalars (Let bind body) = bindsAreScalar && is_scalar scalars' body
+ where
+ (bindsAreScalar, scalars') = is_scalar_bind scalars bind
+ is_scalar scalars (Case e var ty alts)
+ | is_prim_ty ty = is_scalar scalars' e && all (is_scalar_alt scalars') alts
+ | otherwise = False
+ where
+ scalars' = scalars `extendVarSet` var
+ is_scalar scalars (Cast e _coe) = is_scalar scalars e
+ is_scalar scalars (Note _ e ) = is_scalar scalars e
+ is_scalar _scalars (Type {}) = True
+ is_scalar _scalars (Coercion {}) = True
+
+ -- Result: (<is this binding group scalar>, scalars ++ variables bound in this group)
+ is_scalar_bind scalars (NonRec var e) = (is_scalar scalars e, scalars `extendVarSet` var)
+ is_scalar_bind scalars (Rec bnds) = (all (is_scalar scalars') es, scalars')
+ where
+ (vars, es) = unzip bnds
+ scalars' = scalars `extendVarSetList` vars
+
+ is_scalar_alt scalars (_, vars, e) = is_scalar (scalars `extendVarSetList ` vars) e
+
+ -- Checks whether the type might be a parallel array type. In particular, if the outermost
+ -- constructor is a type family, we conservatively assume that it may be a parallel array type.
+ maybe_parr_ty :: Type -> Bool
+ maybe_parr_ty ty
+ | Just ty' <- coreView ty = maybe_parr_ty ty'
+ | Just (tyCon, _) <- splitTyConApp_maybe ty = isPArrTyCon tyCon || isSynFamilyTyCon tyCon
+ maybe_parr_ty _ = False
+
+ -- FIXME: I'm not convinced that this reasoning is (always) sound. If the identify functions
+ -- is called by some other function that is otherwise scalar, it would be very bad
+ -- that just this call to the identity makes it not be scalar.