Allow arbitrary expressions in size expressions.

We still only permit elaboration of expressions that correspond to
variables or integer constants.  This is a step on the path to
realising #1659.

src/Futhark/CLI/Dataset.hs

@@ -198,8 +198,8 @@ toValueType (TEArray d t _) = do
   V.ValueType ds t' <- toValueType t
   pure $ V.ValueType (d' : ds) t'
   where
-    constantDim (SizeExpConst k _) = Right k
-    constantDim _ = Left "Array has non-constant size."
+    constantDim (SizeExp (IntLit k _ _) _) = Right $ fromInteger k
+    constantDim _ = Left "Array has non-constant dimension declaration."
 toValueType (TEVar (QualName [] v) _)
   | Just t <- lookup v m = Right $ V.ValueType [] t
   where

src/Futhark/Doc/Generator.hs

@@ -627,7 +627,7 @@ synopsisSpec spec = case spec of
     specRow (keyword "module " <> vnameSynopsisDef name) ": " <$> synopsisSigExp sig
   IncludeSpec e _ -> fullRow . (keyword "include " <>) <$> synopsisSigExp e
 
-typeExpHtml :: TypeExp VName -> DocM Html
+typeExpHtml :: TypeExp Info VName -> DocM Html
 typeExpHtml e = case e of
   TEUnique t _ -> ("*" <>) <$> typeExpHtml t
   TEArray d at _ -> do
@@ -709,13 +709,12 @@ dimDeclHtml (NamedSize v) = brackets <$> qualNameHtml v
 dimDeclHtml (ConstSize n) = pure $ brackets $ toHtml (show n)
 dimDeclHtml AnySize {} = pure $ brackets mempty
 
-dimExpHtml :: SizeExp VName -> DocM Html
-dimExpHtml SizeExpAny = pure $ brackets mempty
-dimExpHtml (SizeExpNamed v _) = brackets <$> qualNameHtml v
-dimExpHtml (SizeExpConst n _) = pure $ brackets $ toHtml (show n)
+dimExpHtml :: SizeExp Info VName -> DocM Html
+dimExpHtml (SizeExpAny _) = pure $ brackets mempty
+dimExpHtml (SizeExp e _) = pure $ brackets $ toHtml $ prettyString e
 
-typeArgExpHtml :: TypeArgExp VName -> DocM Html
-typeArgExpHtml (TypeArgExpDim d _) = dimExpHtml d
+typeArgExpHtml :: TypeArgExp Info VName -> DocM Html
+typeArgExpHtml (TypeArgExpSize d) = dimExpHtml d
 typeArgExpHtml (TypeArgExpType d) = typeExpHtml d
 
 typeParamHtml :: TypeParam -> Html

src/Futhark/Internalise/Defunctionalise.hs

@@ -143,18 +143,11 @@ replaceStaticValSizes globals orig_substs sv =
         loc
     onExp substs e = onAST substs e
 
-    onTypeExpDim substs d@(SizeExpNamed v loc) =
-      case M.lookup (qualLeaf v) substs of
-        Just (SubstNamed v') ->
-          SizeExpNamed v' loc
-        Just (SubstConst x) ->
-          SizeExpConst x loc
-        Nothing ->
-          d
-    onTypeExpDim _ d = d
+    onTypeExpDim substs (SizeExp e loc) = SizeExp (onExp substs e) loc
+    onTypeExpDim _ (SizeExpAny loc) = SizeExpAny loc
 
-    onTypeArgExp substs (TypeArgExpDim d loc) =
-      TypeArgExpDim (onTypeExpDim substs d) loc
+    onTypeArgExp substs (TypeArgExpSize d) =
+      TypeArgExpSize (onTypeExpDim substs d)
     onTypeArgExp substs (TypeArgExpType te) =
       TypeArgExpType (onTypeExp substs te)
 
@@ -287,7 +280,7 @@ patternArraySizes = arraySizes . patternStructType
 
 data SizeSubst
   = SubstNamed (QualName VName)
-  | SubstConst Int
+  | SubstConst Int64
   deriving (Eq, Ord, Show)
 
 dimMapping ::

src/Futhark/Internalise/Defunctorise.hs

@@ -264,7 +264,7 @@ transformNames x = do
               astMap (substituter $ modScope mod <> scope) e'
         _ -> astMap (substituter scope) e
 
-transformTypeExp :: TypeExp VName -> TransformM (TypeExp VName)
+transformTypeExp :: TypeExp Info VName -> TransformM (TypeExp Info VName)
 transformTypeExp = transformNames
 
 transformStructType :: StructType -> TransformM StructType

src/Futhark/Internalise/Entry.hs

@@ -30,7 +30,7 @@ visibleTypes = VisibleTypes . foldMap (modTypes . snd)
     decTypes (E.TypeDec tb) = [tb]
     decTypes _ = []
 
-findType :: VName -> VisibleTypes -> Maybe (E.TypeExp VName)
+findType :: VName -> VisibleTypes -> Maybe (E.TypeExp E.Info VName)
 findType v (VisibleTypes ts) = E.typeExp <$> find ((== v) . E.typeAlias) ts
 
 valueType :: I.TypeBase I.Rank Uniqueness -> I.ValueType
@@ -39,18 +39,18 @@ valueType (I.Array pt rank _) = I.ValueType I.Signed rank pt
 valueType I.Acc {} = error "valueType Acc"
 valueType I.Mem {} = error "valueType Mem"
 
-withoutDims :: E.TypeExp VName -> (Int, E.TypeExp VName)
+withoutDims :: E.TypeExp E.Info VName -> (Int, E.TypeExp E.Info VName)
 withoutDims (E.TEArray _ te _) =
   let (d, te') = withoutDims te
    in (d + 1, te')
 withoutDims te = (0 :: Int, te)
 
-rootType :: E.TypeExp VName -> E.TypeExp VName
-rootType (E.TEApply te E.TypeArgExpDim {} _) = rootType te
+rootType :: E.TypeExp E.Info VName -> E.TypeExp E.Info VName
+rootType (E.TEApply te E.TypeArgExpSize {} _) = rootType te
 rootType (E.TEUnique te _) = rootType te
 rootType te = te
 
-typeExpOpaqueName :: E.TypeExp VName -> Name
+typeExpOpaqueName :: E.TypeExp E.Info VName -> Name
 typeExpOpaqueName = f . rootType
   where
     f (E.TEArray _ te _) =
@@ -66,7 +66,7 @@ runGenOpaque = flip runState mempty
 addType :: Name -> I.OpaqueType -> GenOpaque ()
 addType s t = modify (<> I.OpaqueTypes [(s, t)])
 
-isRecord :: VisibleTypes -> E.TypeExp VName -> Maybe (M.Map Name (E.TypeExp VName))
+isRecord :: VisibleTypes -> E.TypeExp E.Info VName -> Maybe (M.Map Name (E.TypeExp E.Info VName))
 isRecord _ (E.TERecord fs _) = Just $ M.fromList fs
 isRecord _ (E.TETuple fs _) = Just $ E.tupleFields fs
 isRecord types (E.TEVar v _) = isRecord types =<< findType (E.qualLeaf v) types
@@ -75,7 +75,7 @@ isRecord _ _ = Nothing
 recordFields ::
   VisibleTypes ->
   M.Map Name E.StructType ->
-  Maybe (E.TypeExp VName) ->
+  Maybe (E.TypeExp E.Info VName) ->
   [(Name, E.EntryType)]
 recordFields types fs t =
   case isRecord types . rootType =<< t of

src/Futhark/Internalise/Exps.hs

@@ -2142,7 +2142,7 @@ partitionWithSOACS k lam arrs = do
           (resultBodyM [this_one])
           (resultBodyM [next_one])
 
-typeExpForError :: E.TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
+typeExpForError :: E.TypeExp Info VName -> InternaliseM [ErrorMsgPart SubExp]
 typeExpForError (E.TEVar qn _) =
   pure [ErrorString $ prettyText qn]
 typeExpForError (E.TEUnique te _) =
@@ -2152,10 +2152,8 @@ typeExpForError (E.TEDim dims te _) =
   where
     dims' = mconcat (map onDim dims)
     onDim d = "[" <> prettyText d <> "]"
-typeExpForError (E.TEArray d te _) = do
-  d' <- dimExpForError d
-  te' <- typeExpForError te
-  pure $ ["[", d', "]"] ++ te'
+typeExpForError (E.TEArray d te _) =
+  (<>) <$> sizeExpForError d <*> typeExpForError te
 typeExpForError (E.TETuple tes _) = do
   tes' <- mapM typeExpForError tes
   pure $ ["("] ++ intercalate [", "] tes' ++ [")"]
@@ -2173,7 +2171,7 @@ typeExpForError (E.TEApply t arg _) = do
   t' <- typeExpForError t
   arg' <- case arg of
     TypeArgExpType argt -> typeExpForError argt
-    TypeArgExpDim d _ -> pure <$> dimExpForError d
+    TypeArgExpSize d -> sizeExpForError d
   pure $ t' ++ [" "] ++ arg'
 typeExpForError (E.TESum cs _) = do
   cs' <- mapM (onClause . snd) cs
@@ -2183,16 +2181,11 @@ typeExpForError (E.TESum cs _) = do
       c' <- mapM typeExpForError c
       pure $ intercalate [" "] c'
 
-dimExpForError :: E.SizeExp VName -> InternaliseM (ErrorMsgPart SubExp)
-dimExpForError (SizeExpNamed d _) = do
-  substs <- lookupSubst $ E.qualLeaf d
-  d' <- case substs of
-    Just [v] -> pure v
-    _ -> pure $ I.Var $ E.qualLeaf d
-  pure $ ErrorVal int64 d'
-dimExpForError (SizeExpConst d _) =
-  pure $ ErrorString $ prettyText d
-dimExpForError SizeExpAny = pure ""
+sizeExpForError :: E.SizeExp Info VName -> InternaliseM [ErrorMsgPart SubExp]
+sizeExpForError (SizeExp e _) = do
+  e' <- internaliseExp1 "size" e
+  pure ["[", ErrorVal int64 e', "]"]
+sizeExpForError SizeExpAny {} = pure ["[]"]
 
 -- A smart constructor that compacts neighbouring literals for easier
 -- reading in the IR.

src/Language/Futhark/Parser/Parser.y

@@ -237,7 +237,7 @@ SigExp :: { UncheckedSigExp }
                                 in SigArrow (Just name) $4 $7 (srcspan $1 $>) }
         | SigExp '->' SigExp  { SigArrow Nothing $1 $3 (srcspan $1 $>) }
 
-TypeRef :: { TypeRefBase Name }
+TypeRef :: { TypeRefBase NoInfo Name }
          : QualName TypeParams '=' TypeExpTerm
            { TypeRef (fst $1) $2 $4 (srcspan (snd $1) $>) }
 
@@ -466,18 +466,6 @@ TypeExpTerm :: { UncheckedTypeExp }
          | SumClauses %prec sumprec
            { let (cs, loc) = $1 in TESum cs (srclocOf loc) }
 
-         -- Errors
-         | '[' SizeExp ']' %prec bottom
-           {% parseErrorAt (srcspan $1 $>) $ Just $
-                T.unlines ["Missing array row type.",
-                           "Did you mean []"  <> prettyText $2 <> "?"]
-           }
-         | '...[' SizeExp ']' %prec bottom
-           {% parseErrorAt (srcspan $1 $>) $ Just $
-                T.unlines ["Missing array row type.",
-                           "Did you mean []"  <> prettyText $2 <> "?"]
-           }
-
 SumClauses :: { ([(Name, [UncheckedTypeExp])], Loc) }
             : SumClauses '|' SumClause %prec sumprec
               { let (cs, loc1) = $1; (c, ts, loc2) = $3
@@ -505,20 +493,16 @@ TypeExpAtom :: { UncheckedTypeExp }
              | '(' TypeExp ',' TupleTypes ')' { TETuple ($2:$4) (srcspan $1 $>) }
              | '{' '}'                        { TERecord [] (srcspan $1 $>) }
              | '{' FieldTypes1 '}'            { TERecord $2 (srcspan $1 $>) }
-             | '[' SizeExp ']' TypeExpTerm
-               { TEArray $2 $4 (srcspan $1 $>) }
-             | '...[' SizeExp ']' TypeExpTerm
-               { TEArray $2 $4 (srcspan $1 $>) }
+             | SizeExp TypeExpTerm
+               { TEArray $1 $2 (srcspan $1 $>) }
              | QualName                       { TEVar (fst $1) (srclocOf (snd $1)) }
 
 Constr :: { (Name, Loc) }
         : constructor { let L _ (CONSTRUCTOR c) = $1 in (c, locOf $1) }
 
-TypeArg :: { TypeArgExp Name }
-         : '[' SizeExp ']' %prec top
-           { TypeArgExpDim $2 (srcspan $1 $>) }
-         | '...[' SizeExp ']' %prec top
-           { TypeArgExpDim $2 (srcspan $1 $>) }
+TypeArg :: { TypeArgExp NoInfo Name }
+         : SizeExp %prec top
+           { TypeArgExpSize $1 }
          | TypeExpAtom
            { TypeArgExpType $1 }
 
@@ -533,17 +517,12 @@ TupleTypes :: { [UncheckedTypeExp] }
             : TypeExp                { [$1] }
             | TypeExp ',' TupleTypes { $1 : $3 }
 
-SizeExp :: { SizeExp Name }
-        : QualName
-          { SizeExpNamed (fst $1) (srclocOf (snd $1)) }
-        | intlit
-          { let L loc (INTLIT n) = $1
-            in SizeExpConst (fromIntegral n) (srclocOf loc) }
-        | natlit
-          { let L loc (NATLIT _ n) = $1
-            in SizeExpConst (fromIntegral n) (srclocOf loc) }
-        |
-          { SizeExpAny }
+
+SizeExp :: { SizeExp NoInfo Name }
+         : '[' Exp ']'    { SizeExp $2 (srcspan $1 $>) }
+         | '['     ']'    { SizeExpAny (srcspan $1 $>) }
+         | '...[' Exp ']' { SizeExp $2 (srcspan $1 $>) }
+         | '...['     ']' { SizeExpAny (srcspan $1 $>) }
 
 FunParam :: { PatBase NoInfo Name }
 FunParam : InnerPat { $1 }

src/Language/Futhark/Pretty.hs

@@ -86,10 +86,9 @@ instance Pretty Size where
   pretty (NamedSize v) = pretty v
   pretty (ConstSize n) = pretty n
 
-instance IsName vn => Pretty (SizeExp vn) where
-  pretty SizeExpAny = mempty
-  pretty (SizeExpNamed v _) = pretty v
-  pretty (SizeExpConst n _) = pretty n
+instance (Eq vn, IsName vn, Annot f) => Pretty (SizeExp f vn) where
+  pretty SizeExpAny {} = brackets mempty
+  pretty (SizeExp e _) = brackets $ pretty e
 
 instance Pretty (Shape Size) where
   pretty (Shape ds) = mconcat (map (brackets . pretty) ds)
@@ -168,7 +167,7 @@ prettyTypeArg p (TypeArgType t _) = prettyType p t
 instance Pretty (TypeArg Size) where
   pretty = prettyTypeArg 0
 
-instance (Eq vn, IsName vn) => Pretty (TypeExp vn) where
+instance (Eq vn, IsName vn, Annot f) => Pretty (TypeExp f vn) where
   pretty (TEUnique t _) = "*" <> pretty t
   pretty (TEArray d at _) = brackets (pretty d) <> pretty at
   pretty (TETuple ts _) = parens $ commasep $ map pretty ts
@@ -188,8 +187,8 @@ instance (Eq vn, IsName vn) => Pretty (TypeExp vn) where
   pretty (TEDim dims te _) =
     "?" <> mconcat (map (brackets . prettyName) dims) <> "." <> pretty te
 
-instance (Eq vn, IsName vn) => Pretty (TypeArgExp vn) where
-  pretty (TypeArgExpDim d _) = brackets $ pretty d
+instance (Eq vn, IsName vn, Annot f) => Pretty (TypeArgExp f vn) where
+  pretty (TypeArgExpSize d) = pretty d
   pretty (TypeArgExpType t) = pretty t
 
 instance IsName vn => Pretty (QualName vn) where

src/Language/Futhark/Prop.hs

@@ -1374,7 +1374,7 @@ progHoles = foldMap holesInDec . progDecs
 type UncheckedType = TypeBase (Shape Name) ()
 
 -- | An expression with no type annotations.
-type UncheckedTypeExp = TypeExp Name
+type UncheckedTypeExp = TypeExp NoInfo Name
 
 -- | An identifier with no type annotations.
 type UncheckedIdent = IdentBase NoInfo Name

src/Language/Futhark/Query.hs

@@ -206,7 +206,7 @@ contains a pos =
     Loc start end -> pos >= start && pos <= end
     NoLoc -> False
 
-atPosInTypeExp :: TypeExp VName -> Pos -> Maybe RawAtPos
+atPosInTypeExp :: TypeExp Info VName -> Pos -> Maybe RawAtPos
 atPosInTypeExp te pos =
   case te of
     TEVar qn loc -> do
@@ -229,12 +229,10 @@ atPosInTypeExp te pos =
     TEDim _ t _ ->
       atPosInTypeExp t pos
   where
-    inArg (TypeArgExpDim dim _) = inDim dim
+    inArg (TypeArgExpSize dim) = inDim dim
     inArg (TypeArgExpType e2) = atPosInTypeExp e2 pos
-    inDim (SizeExpNamed qn loc) = do
-      guard $ loc `contains` pos
-      Just $ RawAtName qn $ locOf loc
-    inDim _ = Nothing
+    inDim (SizeExp e _) = atPosInExp e pos
+    inDim SizeExpAny {} = Nothing
 
 atPosInPat :: Pat -> Pos -> Maybe RawAtPos
 atPosInPat (Id vn _ loc) pos = do