github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/gtl/generate_randomized_freepool.py (about)

     1  #!/usr/bin/env python3
     2  
     3  import argparse
     4  import logging
     5  import os
     6  import re
     7  import subprocess
     8  import sys
     9  from typing import List
    10  
    11  
    12  def main() -> None:
    13      logging.basicConfig(level=logging.DEBUG)
    14      output = None
    15      package = None
    16      generate_argv: List[str] = []
    17      for arg in sys.argv[1:]:
    18          m = re.match("^--package=(.*)", arg)
    19          if m:
    20              package = m[1]
    21          m = re.match("^--output=(.*)", arg)
    22          if m:
    23              output = m[1]
    24          else:
    25              generate_argv.append(arg)
    26  
    27      if not output or not package:
    28          raise Exception("--output and --package not set")
    29  
    30      gtl_dir = os.path.dirname(sys.argv[0])
    31      output_dir = os.path.dirname(output)
    32      if not output_dir:
    33          output_dir = "."
    34  
    35      cmdline = (
    36          [sys.executable, os.path.join(gtl_dir, "generate.py"), f"--output={output}.go"]
    37          + generate_argv
    38          + [os.path.join(gtl_dir, "randomized_freepool.go.tpl")]
    39      )
    40      logging.debug("CMDLINE %s", cmdline)
    41      subprocess.check_call(
    42          [sys.executable, os.path.join(gtl_dir, "generate.py"), f"--output={output}.go"]
    43          + generate_argv
    44          + [os.path.join(gtl_dir, "randomized_freepool.go.tpl")]
    45      )
    46      subprocess.check_call(
    47          [
    48              sys.executable,
    49              os.path.join(gtl_dir, "generate.py"),
    50              f"--output={output}_race.go",
    51          ]
    52          + generate_argv
    53          + [os.path.join(gtl_dir, "randomized_freepool_race.go.tpl")]
    54      )
    55  
    56      with open(f"{output_dir}/randomized_freepool_internal.s", "w") as fd:
    57          fd.write(
    58              """// Code generated by generate_randomized_freepool.py. DO NOT EDIT.
    59  
    60  // Dummy file to force the go compiler to honor go:linkname directives. See
    61  //
    62  // https://github.com/golang/go/issues/15006
    63  // http://www.alangpierce.com/blog/2016/03/17/adventures-in-go-accessing-unexported-functions/
    64  """
    65          )
    66  
    67      with open(f"{output_dir}/randomized_freepool_internal.go", "w") as fd:
    68          fd.write(
    69              f"""// Code generated by generate_randomized_freepool.py. DO NOT EDIT.
    70  package {package}
    71  
    72  // This import is needed to use go:linkname.
    73  import _ "unsafe"
    74  
    75  // The following functions are defined in go runtime.  To use them, we need to
    76  // import "unsafe", and elsewhere in this package, import "C" to force compiler
    77  // to recognize the "go:linktime" directive. Some of the details are explained
    78  // in the below blog post.
    79  //
    80  // procPin() pins the caller to the current processor, and returns the processor
    81  // id in range [0,GOMAXPROCS). procUnpin() undos the effect of procPin().
    82  //
    83  // http://www.alangpierce.com/blog/2016/03/17/adventures-in-go-accessing-unexported-functions/
    84  
    85  //go:linkname runtime_procPin sync.runtime_procPin
    86  //go:nosplit
    87  func runtime_procPin() int
    88  
    89  //go:linkname runtime_procUnpin sync.runtime_procUnpin
    90  //go:nosplit
    91  func runtime_procUnpin()
    92  
    93  //go:linkname fastrandn sync.fastrandn
    94  func fastrandn(n uint32) uint32
    95  """
    96          )
    97  
    98  
    99  main()