add variable declaration with binary expressions

src/lang/recast.test.ts

@@ -1,25 +1,40 @@
-import { recast } from "./recast";
-import { Program } from "./abstractSyntaxTree";
-import { abstractSyntaxTree } from "./abstractSyntaxTree";
-import { lexer } from "./tokeniser";
-import { Token } from "./tokeniser";
+import { recast } from './recast'
+import { Program } from './abstractSyntaxTree'
+import { abstractSyntaxTree } from './abstractSyntaxTree'
+import { lexer } from './tokeniser'
+import { Token } from './tokeniser'
 
-describe("recast", () => {
-  it("recasts a simple program", () => {
-    const code = "1 + 2";
-    const {ast, tokens } = code2ast(code);
-    const recasted = recast(ast);
-    expect(recasted).toBe(code);
-  });
-});
+describe('recast', () => {
+  it('recasts a simple program', () => {
+    const code = '1 + 2'
+    const { ast } = code2ast(code)
+    const recasted = recast(ast)
+    expect(recasted).toBe(code)
+  })
+  it('variable declaration', () => {
+    const code = 'const myVar = 5'
+    const { ast } = code2ast(code)
+    const recasted = recast(ast)
+    expect(recasted).toBe(code)
+  })
+  it("variable declaration that's binary with string", () => {
+    const code = "const myVar = 5 + 'yo'"
+    const { ast } = code2ast(code)
+    const recasted = recast(ast)
+    expect(recasted).toBe(code)
+    const codeWithOtherQuotes = 'const myVar = 5 + "yo"'
+    const { ast: ast2 } = code2ast(codeWithOtherQuotes)
+    expect(recast(ast2)).toBe(codeWithOtherQuotes)
+  })
+})
 
 // helpers
 
-function code2ast(code: string): { ast: Program; tokens: Token[]} {
-  const tokens = lexer(code);
-  const ast = abstractSyntaxTree(tokens);
+function code2ast(code: string): { ast: Program; tokens: Token[] } {
+  const tokens = lexer(code)
+  const ast = abstractSyntaxTree(tokens)
   return {
     ast,
-    tokens
+    tokens,
   }
 }

src/lang/recast.ts

@@ -1,27 +1,46 @@
-import { Program, BinaryExpression, BinaryPart } from "./abstractSyntaxTree";
+import { Program, BinaryExpression, BinaryPart } from './abstractSyntaxTree'
 
 export function recast(ast: Program, previousWrittenCode = ''): string {
   return ast.body
     .map((statement) => {
-        if(statement.type === "ExpressionStatement") {
-            if(statement.expression.type === "BinaryExpression") {
-                return recastBinaryExpression(statement.expression);
-            }
+      if (statement.type === 'ExpressionStatement') {
+        if (statement.expression.type === 'BinaryExpression') {
+          return recastBinaryExpression(statement.expression)
         }
-      return statement.type;
+      } else if (statement.type === 'VariableDeclaration') {
+        return statement.declarations
+          .map((declaration) => {
+            if (declaration.init.type === 'BinaryExpression') {
+              return `${statement.kind} ${
+                declaration.id.name
+              } = ${recastBinaryExpression(declaration.init)}`
+            } else if (declaration.init.type === 'Literal') {
+              return `${statement.kind} ${declaration.id.name} = ${declaration.init.value}`
+            }
+            return ''
+          })
+          .join('')
+      }
+      return statement.type
     })
-    .join("\n");
+    .join('\n')
 }
 
 function recastBinaryExpression(expression: BinaryExpression): string {
-    return `${recastBinaryPart(expression.left)} ${expression.operator} ${recastBinaryPart(expression.right)}`
+  return `${recastBinaryPart(expression.left)} ${
+    expression.operator
+  } ${recastBinaryPart(expression.right)}`
 }
 
 function recastBinaryPart(part: BinaryPart): string {
-    if(part.type === "Literal") {
-        return String(part?.value);
-    } else if(part.type === "Identifier") {
-        return part.name;
+  if (part.type === 'Literal') {
+    if (typeof part.value === 'string') {
+      const quote = part.raw.includes('"') ? '"' : "'"
+      return `${quote}${part.value}${quote}`
     }
-    throw new Error(`Cannot recast ${part}`);
+    return String(part?.value)
+  } else if (part.type === 'Identifier') {
+    return part.name
+  }
+  throw new Error(`Cannot recast ${part}`)
 }