github.com/qri-io/qri@v0.10.1-0.20220104210721-c771715036cb/transform/staticlark/unit.go (about)

     1  package staticlark
     2  
     3  import (
     4  	"fmt"
     5  	"strings"
     6  	"unicode"
     7  
     8  	"go.starlark.net/syntax"
     9  )
    10  
    11  // unit of code, roughly equivalent to a single line
    12  // can be an assignment, function call, or control structure such as `if` or `for`
    13  // represented as a lisp-like AST
    14  type unit struct {
    15  	atom  string
    16  	tail  []*unit
    17  	todo  bool
    18  	where syntax.Position
    19  }
    20  
    21  // String converts a unit into a prefix notation string
    22  func (u *unit) String() string {
    23  	// atom only
    24  	if u.tail == nil {
    25  		if u.todo {
    26  			return fmt.Sprintf("TODO:%s", u.atom)
    27  		}
    28  		return u.atom
    29  	}
    30  	// function with no args
    31  	if len(u.tail) == 0 {
    32  		if u.todo {
    33  			return fmt.Sprintf("TODO:[%s]", u.atom)
    34  		}
    35  		return fmt.Sprintf("[%s]", u.atom)
    36  	}
    37  	// function with args
    38  	res := make([]string, len(u.tail))
    39  	for i, node := range u.tail {
    40  		res[i] = node.String()
    41  	}
    42  	if u.todo {
    43  		return fmt.Sprintf("TODO:[%s %s]", u.atom, strings.Join(res, " "))
    44  	}
    45  	return fmt.Sprintf("[%s %s]", u.atom, strings.Join(res, " "))
    46  }
    47  
    48  // Push adds to the tail of the unit
    49  func (u *unit) Push(text string) {
    50  	u.tail = append(u.tail, &unit{atom: text})
    51  }
    52  
    53  // return all of the sources of data, as lexical tokens, that affect
    54  // the calculation of this unit's value
    55  func (u *unit) DataSources() []string {
    56  	startIndex := 1
    57  	if u.atom == "set!" {
    58  		startIndex = 2
    59  	}
    60  	if len(u.tail) < startIndex {
    61  		return []string{}
    62  	}
    63  	res := []string{}
    64  	for i, t := range u.tail {
    65  		if i < (startIndex - 1) {
    66  			continue
    67  		}
    68  		res = append(res, t.getSources()...)
    69  	}
    70  	return res
    71  }
    72  
    73  func (u *unit) getSources() []string {
    74  	if len(u.tail) == 0 {
    75  		if isIdent(u.atom) {
    76  			return []string{u.atom}
    77  		}
    78  		return []string{}
    79  	}
    80  	res := []string{}
    81  	for _, t := range u.tail {
    82  		res = append(res, t.getSources()...)
    83  	}
    84  	return res
    85  }
    86  
    87  type invoke struct {
    88  	Name string   `json:"name"`
    89  	Args []string `json:"args"`
    90  }
    91  
    92  func (u *unit) Invocations() []invoke {
    93  	if u.atom == "set!" {
    94  		return u.tail[1].Invocations()
    95  	}
    96  	if u.atom == "return" {
    97  		return u.invokeFromList()
    98  	}
    99  	if isIdent(u.atom) && u.tail != nil {
   100  		name := u.atom
   101  		args := u.getSources()
   102  		inv := invoke{Name: name, Args: args}
   103  		return append([]invoke{inv}, u.invokeFromList()...)
   104  	}
   105  	return u.invokeFromList()
   106  }
   107  
   108  func (u *unit) invokeFromList() []invoke {
   109  	if u.tail == nil {
   110  		return []invoke{}
   111  	}
   112  	var res []invoke
   113  	for _, t := range u.tail {
   114  		res = append(res, t.Invocations()...)
   115  	}
   116  	if res == nil {
   117  		return []invoke{}
   118  	}
   119  	return res
   120  }
   121  
   122  func (u *unit) AssignsTo() string {
   123  	if u.atom == "set!" {
   124  		return u.tail[0].atom
   125  	}
   126  	return ""
   127  }
   128  
   129  func (u *unit) IsReturn() bool {
   130  	return u.atom == "return"
   131  }
   132  
   133  func toUnitTODO(text string) *unit {
   134  	return &unit{atom: text, todo: true}
   135  }
   136  
   137  func isIdent(text string) bool {
   138  	return unicode.IsLetter([]rune(text)[0])
   139  }