github.com/coyove/nj@v0.0.0-20221110084952-c7f8db1065c3/cmd/script/main.go (about)

     1  package main
     2  
     3  import (
     4  	"bufio"
     5  	"flag"
     6  	"fmt"
     7  	"io/ioutil"
     8  	"log"
     9  	"net/http"
    10  	"os"
    11  	"runtime"
    12  	"runtime/pprof"
    13  	"strconv"
    14  	"strings"
    15  	"time"
    16  
    17  	"github.com/coyove/nj"
    18  	"github.com/coyove/nj/bas"
    19  )
    20  
    21  var (
    22  	output    = flag.String("o", "ret", "separated by comma: (none|compileonly|opcode|bytes|ret|timing)+")
    23  	input     = flag.String("i", "f", "input source, 'f': file, '-': stdin, others: string")
    24  	version   = flag.Bool("v", false, "print version and usage")
    25  	repl      = flag.Bool("repl", false, "repl mode")
    26  	timeout   = flag.Int("t", 0, "max execution time in ms")
    27  	stackSize = flag.Int("ss", 1e6, "max stack size, 1 stack size is 16-byte on 64bit platform")
    28  	apiServer = flag.String("serve", "", "start as language playground")
    29  )
    30  
    31  func main() {
    32  	source := ""
    33  	for i, arg := range os.Args {
    34  		if _, err := os.Stat(arg); err == nil && i > 0 {
    35  			source = arg
    36  			os.Args = append(os.Args[:i], os.Args[i+1:]...)
    37  			break
    38  		}
    39  	}
    40  
    41  	flag.Parse()
    42  
    43  	if *apiServer != "" {
    44  		http.HandleFunc("/", nj.PlaygroundHandler("", nil))
    45  		log.Println("listen", *apiServer)
    46  		http.ListenAndServe(*apiServer, nil)
    47  		return
    48  	}
    49  
    50  	if *repl {
    51  		runRepl()
    52  	}
    53  
    54  	log.SetFlags(0)
    55  
    56  	if *version {
    57  		fmt.Println("nj virtual machine v" + strconv.FormatInt(bas.Version, 10) + " (" + runtime.GOOS + "/" + runtime.GOARCH + ")")
    58  		flag.Usage()
    59  		return
    60  	}
    61  
    62  	if p, _ := os.Getwd(); !strings.Contains(p, "/cmd/nj") {
    63  		f, err := os.Create("cpuprofile")
    64  		if err != nil {
    65  			log.Fatal(err)
    66  		}
    67  		pprof.StartCPUProfile(f)
    68  		defer pprof.StopCPUProfile()
    69  	}
    70  
    71  	switch *input {
    72  	case "f":
    73  		if source == "" {
    74  			log.Fatalln("Please specify the input file: ./nj <filename>")
    75  		}
    76  	case "-":
    77  		buf, err := ioutil.ReadAll(os.Stdin)
    78  		if err != nil {
    79  			log.Fatalln(err)
    80  		}
    81  		source = string(buf)
    82  	default:
    83  		if _, err := os.Stat(*input); err == nil {
    84  			source = *input
    85  			*input = "f"
    86  		} else {
    87  			source = *input
    88  		}
    89  	}
    90  
    91  	var _opcode, _timing, _ret, _compileonly bool
    92  	for _, a := range strings.Split(*output, ",") {
    93  		switch a {
    94  		case "o", "opcode", "op":
    95  			_opcode = true
    96  		case "r", "ret", "return":
    97  			_ret = true
    98  		case "t", "time", "timing":
    99  			_timing = true
   100  		case "co", "compile", "compileonly":
   101  			_compileonly = true
   102  		}
   103  	}
   104  
   105  	runtime.GOMAXPROCS(runtime.NumCPU() * 2)
   106  	start := time.Now()
   107  
   108  	var b *bas.Program
   109  	var err error
   110  
   111  	defer func() {
   112  		if _opcode {
   113  			log.Println(b.GoString())
   114  		}
   115  		if _timing {
   116  			log.Printf("Time elapsed: %v\n", time.Since(start))
   117  		}
   118  	}()
   119  
   120  	if *input == "f" {
   121  		b, err = nj.LoadFile(source, nil)
   122  	} else {
   123  		b, err = nj.LoadString(source, nil)
   124  	}
   125  	if err != nil {
   126  		log.Fatalln(err)
   127  	}
   128  	b.MaxStackSize = int64(*stackSize)
   129  
   130  	if _compileonly {
   131  		return
   132  	}
   133  
   134  	if *timeout > 0 {
   135  		go func() {
   136  			time.Sleep(time.Millisecond * time.Duration(*timeout))
   137  			b.Stop()
   138  		}()
   139  	}
   140  
   141  	i := b.Run()
   142  	if _ret {
   143  		fmt.Println(i)
   144  	}
   145  }
   146  
   147  func runRepl() {
   148  	var code []string
   149  	var globals *bas.Object
   150  	rd := bufio.NewReader(os.Stdin)
   151  	for {
   152  		fmt.Print("> ")
   153  		s, err := rd.ReadString('\n')
   154  		if err != nil {
   155  			fmt.Println("Exit")
   156  			break
   157  		}
   158  		s = strings.TrimSuffix(s, "\n")
   159  		code = append(code, s)
   160  		if s == "" || strings.HasSuffix(s, ";") {
   161  			text := strings.TrimSuffix(strings.Join(code, "\n"), ";")
   162  			p, err := nj.LoadString(text, &nj.LoadOptions{Globals: globals.ToMap()})
   163  			if err != nil {
   164  				fmt.Println("x", err)
   165  			} else {
   166  				res := p.Run()
   167  				fmt.Println("=>", res)
   168  				globals = p.LocalsObject()
   169  			}
   170  			code = code[:0]
   171  		}
   172  	}
   173  }