5 files changed, 23 insertions, 8 deletions

diff --git a/cmd/vm/main.go b/cmd/vm/main.go
index bb8f142..d58dae1 100644
--- a/cmd/vm/main.go
+++ b/cmd/vm/main.go
@@ -15,6 +15,7 @@ func main() {
 	// Take the first argument as the path to the file to be parsed.
 	compile := flag.String("c", "", "compile to file (if not bytecode)")
 	decompile := flag.String("d", "", "decompile bytecode to file")
+	generatePCs := flag.Bool("p", false, "if decompiling, generate PC information")
 	run := flag.Bool("r", false, "run")
 	bytecode := flag.Bool("b", false, "interpret bytecode, without this flag the program will be interpreted in Lang VM text format.")
 
@@ -36,6 +37,10 @@ func main() {
 		exit("nothing to do, either -c, -d or -r is required")
 	}
 
+	if *decompile == "" && *generatePCs {
+		exit("-p requires -d")
+	}
+
 	path := flag.Arg(0)
 	if path == "" {
 		exit("no file specified")
@@ -75,7 +80,7 @@ func main() {
 		bc = code.New(content, code.NewDebugInfo(path))
 
 		if *decompile != "" {
-			decomp := text.NewDecompiler(bc)
+			decomp := text.NewDecompiler(bc, *generatePCs)
 			result := decomp.Decompile()
 
 			output, err := os.Create(*decompile)
diff --git a/pkg/bot/cmds.go b/pkg/bot/cmds.go
index 7d451a9..fd3548e 100644
--- a/pkg/bot/cmds.go
+++ b/pkg/bot/cmds.go
@@ -79,7 +79,7 @@ func langCmd(b *Bot, content string, msg events.Message) error {
 
 	b.client.SendMessage(msg.ChannelID, "successfully compiled code")
 
-	decomp := text.NewDecompiler(*module.Code())
+	decomp := text.NewDecompiler(*module.Code(), false)
 	decompRes := decomp.Decompile()
 
 	b.client.SendMessage(msg.ChannelID, fmt.Sprintf("resulting decompiled bytecode:\n```\n%s\n```", decompRes))
diff --git a/pkg/lang/compiler/compiler_test.go b/pkg/lang/compiler/compiler_test.go
index 2e4e6e7..533bb43 100644
--- a/pkg/lang/compiler/compiler_test.go
+++ b/pkg/lang/compiler/compiler_test.go
@@ -687,8 +687,8 @@ func mustCompileTo(t *testing.T, src, expected string) {
 	require.NoError(t, err)
 
 	if !assert.Equal(t, expectedResult.Code(), testResult.Code().Code()) {
-		d1 := text.NewDecompiler(expectedResult)
-		d2 := text.NewDecompiler(*testResult.Code())
+		d1 := text.NewDecompiler(expectedResult, false)
+		d2 := text.NewDecompiler(*testResult.Code(), false)
 		require.Equal(t, d1.Decompile(), d2.Decompile())
 	}
 }
diff --git a/pkg/lang/vm/text/decompiler.go b/pkg/lang/vm/text/decompiler.go
index df9e1ed..ed48412 100644
--- a/pkg/lang/vm/text/decompiler.go
+++ b/pkg/lang/vm/text/decompiler.go
@@ -14,12 +14,15 @@ type Decompiler struct {
 	c code.Code
 
 	pcToLine rangemap.RangeMap[int]
+
+	generatePCs bool
 }
 
-func NewDecompiler(c code.Code) *Decompiler {
+func NewDecompiler(c code.Code, generatePCs bool) *Decompiler {
 	return &Decompiler{
-		c:        c,
-		pcToLine: rangemap.New[int](),
+		c:           c,
+		pcToLine:    rangemap.New[int](),
+		generatePCs: generatePCs,
 	}
 }
 
@@ -27,12 +30,19 @@ func (d *Decompiler) Decompile() string {
 	lines := make([]string, 0)
 	bc := d.c.Code()
 
+	longestPcStringLen := len(strconv.FormatInt(int64(len(bc)), 10))
+
 	for len(bc) != 0 {
+		startingPc := len(d.c.Code()) - len(bc)
+
 		line, rest := d.decompileInstruction(bc)
 		bc = rest
 
 		d.pcToLine.AppendToLast(d.c.Len()-len(bc), len(lines))
 
+		if d.generatePCs {
+			line = fmt.Sprintf("%*d: %s", longestPcStringLen, startingPc, line)
+		}
 		lines = append(lines, line)
 	}
 
diff --git a/pkg/lang/vm/text/decompiler_test.go b/pkg/lang/vm/text/decompiler_test.go
index 01e49b0..5b9a0d7 100644
--- a/pkg/lang/vm/text/decompiler_test.go
+++ b/pkg/lang/vm/text/decompiler_test.go
@@ -69,7 +69,7 @@ func test(t *testing.T, code string, expected string) {
 	resCompiled, err := comp.Compile()
 	require.NoError(t, err)
 
-	decomp := text.NewDecompiler(resCompiled)
+	decomp := text.NewDecompiler(resCompiled, false)
 	resDecompiled := decomp.Decompile()
 	require.Equal(t, trimmedExpected, resDecompiled)
 }