From 878043c39cd60ed0848f4ff8d7b0885580588aed Mon Sep 17 00:00:00 2001
From: Wouter van Oortmerssen <aardappel@gmail.com>
Date: Fri, 28 Jun 2019 20:29:16 +0000
Subject: [PATCH] [WebAssembly] Assembler: Improve section parsing.

Reviewers: sbc100

Subscribers: dschuff, jgravelle-google, aheejin, sunfish, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D63947

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@364681 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/MC/MCParser/WasmAsmParser.cpp             | 43 ++++++-------------
 .../AsmParser/WebAssemblyAsmParser.cpp        |  4 ++
 2 files changed, 18 insertions(+), 29 deletions(-)

diff --git a/lib/MC/MCParser/WasmAsmParser.cpp b/lib/MC/MCParser/WasmAsmParser.cpp
index 3000a2315d3..141e7fc36f9 100644
--- a/lib/MC/MCParser/WasmAsmParser.cpp
+++ b/lib/MC/MCParser/WasmAsmParser.cpp
@@ -114,13 +114,17 @@ public:
     if (Lexer->isNot(AsmToken::String))
       return error("expected string in directive, instead got: ", Lexer->getTok());
 
-    SectionKind Kind = StringSwitch<SectionKind>(Name)
-                       .StartsWith(".data", SectionKind::getData())
-                       .StartsWith(".rodata", SectionKind::getReadOnly())
-                       .StartsWith(".text", SectionKind::getText())
-                       .StartsWith(".custom_section", SectionKind::getMetadata());
-
-    MCSectionWasm* Section = getContext().getWasmSection(Name, Kind);
+    auto Kind = StringSwitch<Optional<SectionKind>>(Name)
+                    .StartsWith(".data", SectionKind::getData())
+                    .StartsWith(".rodata", SectionKind::getReadOnly())
+                    .StartsWith(".text", SectionKind::getText())
+                    .StartsWith(".custom_section", SectionKind::getMetadata())
+                    .StartsWith(".bss", SectionKind::getBSS())
+                    .Default(Optional<SectionKind>());
+    if (!Kind.hasValue())
+      return Parser->Error(Lexer->getLoc(), "unknown section kind: " + Name);
+
+    MCSectionWasm *Section = getContext().getWasmSection(Name, Kind.getValue());
 
     // Update section flags if present in this .section directive
     bool Passive = false;
@@ -139,28 +143,9 @@ public:
     if (expect(AsmToken::Comma, ",") || expect(AsmToken::At, "@") ||
         expect(AsmToken::EndOfStatement, "eol"))
       return true;
-    struct SectionType {
-      const char *Name;
-      SectionKind Kind;
-    };
-    static SectionType SectionTypes[] = {
-        {".text", SectionKind::getText()},
-        {".rodata", SectionKind::getReadOnly()},
-        {".data", SectionKind::getData()},
-        {".custom_section", SectionKind::getMetadata()},
-        // TODO: add more types.
-    };
-    for (size_t I = 0; I < sizeof(SectionTypes) / sizeof(SectionType); I++) {
-      if (Name.startswith(SectionTypes[I].Name)) {
-        auto WS = getContext().getWasmSection(Name, SectionTypes[I].Kind);
-        getStreamer().SwitchSection(WS);
-        return false;
-      }
-    }
-    // Not found, just ignore this section.
-    // For code in a text section WebAssemblyAsmParser automatically adds
-    // one section per function, so they're optional to be specified with
-    // this directive.
+
+    auto WS = getContext().getWasmSection(Name, Kind.getValue());
+    getStreamer().SwitchSection(WS);
     return false;
   }
 
diff --git a/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmParser.cpp b/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmParser.cpp
index 196ef2d9b91..b891f8868ac 100644
--- a/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmParser.cpp
+++ b/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmParser.cpp
@@ -748,6 +748,10 @@ public:
     auto SymName = Symbol->getName();
     if (SymName.startswith(".L"))
       return; // Local Symbol.
+    // Only create a new text section if we're already in one.
+    auto CWS = cast<MCSectionWasm>(getStreamer().getCurrentSection().first);
+    if (!CWS || !CWS->getKind().isText())
+      return;
     auto SecName = ".text." + SymName;
     auto WS = getContext().getWasmSection(SecName, SectionKind::getText());
     getStreamer().SwitchSection(WS);
-- 
2.40.0