From 932f636d8452fa2a67abf0c86ef2abfc71fd2473 Mon Sep 17 00:00:00 2001
From: Steven <stevenlgtm@gmail.com>
Date: Sat, 27 Jan 2024 19:09:10 +0800
Subject: [PATCH] chore: update codeblock renderer

---
 plugin/gomark/parser/code_block.go           | 65 +++++++++++---------
 plugin/gomark/parser/math_block.go           | 13 ++--
 web/src/components/MemoContent/CodeBlock.tsx | 11 +---
 3 files changed, 47 insertions(+), 42 deletions(-)

diff --git a/plugin/gomark/parser/code_block.go b/plugin/gomark/parser/code_block.go
index 9af35998..36eb93f6 100644
--- a/plugin/gomark/parser/code_block.go
+++ b/plugin/gomark/parser/code_block.go
@@ -1,6 +1,8 @@
 package parser
 
 import (
+	"slices"
+
 	"github.com/usememos/memos/plugin/gomark/ast"
 	"github.com/usememos/memos/plugin/gomark/parser/tokenizer"
 )
@@ -15,51 +17,56 @@ func NewCodeBlockParser() *CodeBlockParser {
 }
 
 func (*CodeBlockParser) Match(tokens []*tokenizer.Token) (ast.Node, int) {
-	if len(tokens) < 9 {
+	rows := tokenizer.Split(tokens, tokenizer.Newline)
+	if len(rows) < 3 {
 		return nil, 0
 	}
 
-	if tokens[0].Type != tokenizer.Backtick || tokens[1].Type != tokenizer.Backtick || tokens[2].Type != tokenizer.Backtick {
+	firstRow := rows[0]
+	if len(firstRow) < 3 {
 		return nil, 0
 	}
-	if tokens[3].Type != tokenizer.Newline && tokens[4].Type != tokenizer.Newline {
+	if firstRow[0].Type != tokenizer.Backtick || firstRow[1].Type != tokenizer.Backtick || firstRow[2].Type != tokenizer.Backtick {
 		return nil, 0
 	}
-	cursor := 4
-	if tokens[3].Type != tokenizer.Newline {
-		cursor = 5
+	languageTokens := []*tokenizer.Token{}
+	if len(firstRow) > 3 {
+		languageTokens = firstRow[3:]
+		// Check if language is valid.
+		availableLanguageTokenTypes := []tokenizer.TokenType{tokenizer.Text, tokenizer.Number, tokenizer.Underscore}
+		for _, token := range languageTokens {
+			if !slices.Contains(availableLanguageTokenTypes, token.Type) {
+				return nil, 0
+			}
+		}
 	}
 
+	contentRows := [][]*tokenizer.Token{}
 	matched := false
-	for ; cursor < len(tokens)-3; cursor++ {
-		if tokens[cursor].Type == tokenizer.Newline && tokens[cursor+1].Type == tokenizer.Backtick && tokens[cursor+2].Type == tokenizer.Backtick && tokens[cursor+3].Type == tokenizer.Backtick {
-			if cursor+3 == len(tokens)-1 {
-				cursor += 4
-				matched = true
-				break
-			} else if tokens[cursor+4].Type == tokenizer.Newline {
-				cursor += 4
-				matched = true
-				break
-			}
+	for _, row := range rows[1:] {
+		if len(row) == 3 && row[0].Type == tokenizer.Backtick && row[1].Type == tokenizer.Backtick && row[2].Type == tokenizer.Backtick {
+			matched = true
+			break
 		}
+		contentRows = append(contentRows, row)
 	}
 	if !matched {
 		return nil, 0
 	}
 
-	languageToken := tokens[3]
-	contentStart, contentEnd := 5, cursor-4
-	if languageToken.Type == tokenizer.Newline {
-		languageToken = nil
-		contentStart = 4
+	contentTokens := []*tokenizer.Token{}
+	for index, row := range contentRows {
+		contentTokens = append(contentTokens, row...)
+		if index != len(contentRows)-1 {
+			contentTokens = append(contentTokens, &tokenizer.Token{
+				Type:  tokenizer.Newline,
+				Value: "\n",
+			})
+		}
 	}
 
-	codeBlock := &ast.CodeBlock{
-		Content: tokenizer.Stringify(tokens[contentStart:contentEnd]),
-	}
-	if languageToken != nil {
-		codeBlock.Language = languageToken.String()
-	}
-	return codeBlock, cursor
+	return &ast.CodeBlock{
+		Content:  tokenizer.Stringify(contentTokens),
+		Language: tokenizer.Stringify(languageTokens),
+	}, 4 + len(languageTokens) + len(contentTokens) + 4
 }
diff --git a/plugin/gomark/parser/math_block.go b/plugin/gomark/parser/math_block.go
index 4fd9cc4b..77fe080e 100644
--- a/plugin/gomark/parser/math_block.go
+++ b/plugin/gomark/parser/math_block.go
@@ -38,13 +38,16 @@ func (*MathBlockParser) Match(tokens []*tokenizer.Token) (ast.Node, int) {
 	}
 
 	contentTokens := []*tokenizer.Token{}
-	for _, row := range contentRows {
+	for index, row := range contentRows {
 		contentTokens = append(contentTokens, row...)
-		contentTokens = append(contentTokens, &tokenizer.Token{
-			Type: tokenizer.Newline,
-		})
+		if index != len(contentRows)-1 {
+			contentTokens = append(contentTokens, &tokenizer.Token{
+				Type:  tokenizer.Newline,
+				Value: "\n",
+			})
+		}
 	}
 	return &ast.MathBlock{
 		Content: tokenizer.Stringify(contentTokens),
-	}, 3 + len(contentTokens) + 2
+	}, 3 + len(contentTokens) + 3
 }
diff --git a/web/src/components/MemoContent/CodeBlock.tsx b/web/src/components/MemoContent/CodeBlock.tsx
index 4254b8ab..fbdd7114 100644
--- a/web/src/components/MemoContent/CodeBlock.tsx
+++ b/web/src/components/MemoContent/CodeBlock.tsx
@@ -15,14 +15,9 @@ const CodeBlock: React.FC<Props> = ({ language, content }: Props) => {
   const formatedLanguage = language.toLowerCase() || "plaintext";
   let highlightedCode = hljs.highlightAuto(content).value;
 
-  // Users can set Markdown code blocks as 'iframe'
-  // to embed videos or external audio links from services like Apple Music or Spotify.
-  if (formatedLanguage === "iframe") {
-    const renderHTML = () => {
-      return { __html: content };
-    };
-
-    return <div className="mx-auto w-11/12 !my-4" dangerouslySetInnerHTML={renderHTML()} />;
+  // Users can set Markdown code blocks as `__html` to render HTML directly.
+  if (formatedLanguage === "__html") {
+    return <div className="w-full !my-2" dangerouslySetInnerHTML={{ __html: content }} />;
   }
 
   try {