ssa_generics.mx raw

   1  package ssa
   2  
   3  import (
   4  	"git.smesh.lol/moxie/pkg/token"
   5  	"git.smesh.lol/moxie/pkg/syntax"
   6  	"git.smesh.lol/moxie/pkg/types"
   7  )
   8  
   9  func mangleGenericName(baseName string, typeArgs []types.Type) (s string) {
  10  	s = baseName | "__"
  11  	for i, t := range typeArgs {
  12  		if i > 0 {
  13  			s = s | "_"
  14  		}
  15  		if t == nil {
  16  			s = s | "nil"
  17  		} else {
  18  			ts := t.String()
  19  			for j := 0; j < len(ts); j++ {
  20  				c := ts[j]
  21  				if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') {
  22  					s = s | string([]byte{c})
  23  				} else if c == '.' {
  24  					s = s | "."
  25  				} else {
  26  					s = s | "_"
  27  				}
  28  			}
  29  		}
  30  	}
  31  	return s
  32  }
  33  
  34  func resolveTypeExprWithSubst(e syntax.Expr, typeSubst map[string]types.Type, pkgScope *types.Scope) (t types.Type) {
  35  	if e == nil {
  36  		return nil
  37  	}
  38  	switch e := e.(type) {
  39  	case *syntax.Name:
  40  		if t, ok := typeSubst[e.Value]; ok {
  41  			return t
  42  		}
  43  		_, obj := types.Universe.LookupParent(e.Value)
  44  		if obj != nil {
  45  			if tn, ok := obj.(*types.TypeName); ok {
  46  				return tn.typ
  47  			}
  48  		}
  49  		if pkgScope != nil {
  50  			tobj := pkgScope.Lookup(e.Value)
  51  			if tobj != nil {
  52  				if tn, ok := tobj.(*types.TypeName); ok {
  53  					return tn.typ
  54  				}
  55  			}
  56  		}
  57  		if importRegistry != nil {
  58  			var regKeys []string
  59  			for k := range importRegistry {
  60  				regKeys = append(regKeys, k)
  61  			}
  62  			for i := 1; i < len(regKeys); i++ {
  63  				for j := i; j > 0 && regKeys[j] < regKeys[j-1]; j-- {
  64  					regKeys[j], regKeys[j-1] = regKeys[j-1], regKeys[j]
  65  				}
  66  			}
  67  			for _, k := range regKeys {
  68  				pkg := importRegistry[k]
  69  				tobj := pkg.Scope().Lookup(e.Value)
  70  				if tobj != nil {
  71  					if tn, ok := tobj.(*types.TypeName); ok {
  72  						return tn.typ
  73  					}
  74  				}
  75  			}
  76  		}
  77  		return nil
  78  	case *syntax.DotsType:
  79  		elem := resolveTypeExprWithSubst(e.Elem, typeSubst, pkgScope)
  80  		if elem != nil {
  81  			if b, ok := elem.(*types.Basic); ok && b.Kind() == types.Uint8 {
  82  				return types.Typ[types.TCString]
  83  			}
  84  			return types.NewSlice(elem)
  85  		}
  86  	case *syntax.SliceType:
  87  		elem := resolveTypeExprWithSubst(e.Elem, typeSubst, pkgScope)
  88  		if elem != nil {
  89  			if b, ok := elem.(*types.Basic); ok && b.Kind() == types.Uint8 {
  90  				return types.Typ[types.TCString]
  91  			}
  92  			return types.NewSlice(elem)
  93  		}
  94  	case *syntax.Operation:
  95  		if e.Y == nil && e.Op == token.Mul {
  96  			base := resolveTypeExprWithSubst(e.X, typeSubst, pkgScope)
  97  			if base != nil {
  98  				return types.NewPointer(base)
  99  			}
 100  		}
 101  	case *syntax.MapType:
 102  		key := resolveTypeExprWithSubst(e.Key, typeSubst, pkgScope)
 103  		val := resolveTypeExprWithSubst(e.Value, typeSubst, pkgScope)
 104  		if key != nil && val != nil {
 105  			return types.NewTCMap(key, val)
 106  		}
 107  	case *syntax.ChanType:
 108  		elem := resolveTypeExprWithSubst(e.Elem, typeSubst, pkgScope)
 109  		if elem != nil {
 110  			dir := types.TCSendRecv
 111  			if e.Dir == syntax.SendOnly {
 112  				dir = types.TCSendOnly
 113  			} else if e.Dir == syntax.RecvOnly {
 114  				dir = types.TCRecvOnly
 115  			}
 116  			return types.NewTCChan(dir, elem)
 117  		}
 118  	case *syntax.FuncType:
 119  		return resolveSignatureWithSubst(e, typeSubst, pkgScope)
 120  	case *syntax.SelectorExpr:
 121  		if pkgName, ok := e.X.(*syntax.Name); ok {
 122  			if importRegistry != nil {
 123  				var regKeys []string
 124  				for k := range importRegistry {
 125  					regKeys = append(regKeys, k)
 126  				}
 127  				for i := 1; i < len(regKeys); i++ {
 128  					for j := i; j > 0 && regKeys[j] < regKeys[j-1]; j-- {
 129  						regKeys[j], regKeys[j-1] = regKeys[j-1], regKeys[j]
 130  					}
 131  				}
 132  				for _, k := range regKeys {
 133  					pkg := importRegistry[k]
 134  					if pkg.Name() == pkgName.Value {
 135  						typeObj := pkg.Scope().Lookup(e.Sel.Value)
 136  						if typeObj != nil {
 137  							if tn, ok2 := typeObj.(*types.TypeName); ok2 {
 138  								return tn.typ
 139  							}
 140  						}
 141  						break
 142  					}
 143  				}
 144  			}
 145  		}
 146  	}
 147  	return nil
 148  }
 149  
 150  func resolveSignatureWithSubst(ft *syntax.FuncType, typeSubst map[string]types.Type, pkgScope *types.Scope) (s *types.Signature) {
 151  	if ft == nil {
 152  		return nil
 153  	}
 154  	var params []*types.TCVar
 155  	variadic := false
 156  	for _, p := range ft.ParamList {
 157  		typ := resolveTypeExprWithSubst(p.Type, typeSubst, pkgScope)
 158  		name := ""
 159  		if p.Name != nil {
 160  			name = p.Name.Value
 161  		}
 162  		params = append(params, types.NewTCVar(nil, name, typ))
 163  	}
 164  	if len(ft.ParamList) > 0 {
 165  		last := ft.ParamList[len(ft.ParamList)-1]
 166  		if _, ok := last.Type.(*syntax.DotsType); ok {
 167  			variadic = true
 168  		}
 169  	}
 170  	var results []*types.TCVar
 171  	for _, r := range ft.ResultList {
 172  		typ := resolveTypeExprWithSubst(r.Type, typeSubst, pkgScope)
 173  		name := ""
 174  		if r.Name != nil {
 175  			name = r.Name.Value
 176  		}
 177  		if typ == nil && r.Name != nil {
 178  			rname := ""
 179  			if rn, ok := r.Type.(*syntax.Name); ok {
 180  				rname = rn.Value
 181  			}
 182  		}
 183  		results = append(results, types.NewTCVar(nil, name, typ))
 184  	}
 185  	var paramTuple *types.Tuple
 186  	if len(params) > 0 {
 187  		paramTuple = types.NewTuple(params...)
 188  	}
 189  	var resultTuple *types.Tuple
 190  	if len(results) > 0 {
 191  		resultTuple = types.NewTuple(results...)
 192  	}
 193  	return types.NewSignature(nil, paramTuple, resultTuple, variadic)
 194  }
 195  
 196  func matchTypeParam(paramType syntax.Expr, argType types.Type, tparams []*syntax.Field, typeMap map[string]types.Type) {
 197  	if paramType == nil || argType == nil {
 198  		return
 199  	}
 200  	switch pt := paramType.(type) {
 201  	case *syntax.Name:
 202  		for _, tp := range tparams {
 203  			if tp.Name != nil && tp.Name.Value == pt.Value {
 204  				typeMap[pt.Value] = argType
 205  				return
 206  			}
 207  		}
 208  	case *syntax.DotsType:
 209  		matchTypeParam(pt.Elem, argType, tparams, typeMap)
 210  	case *syntax.SliceType:
 211  		if sl, ok := types.SafeUnderlying(argType).(*types.Slice); ok {
 212  			matchTypeParam(pt.Elem, sl.Elem(), tparams, typeMap)
 213  		}
 214  	case *syntax.FuncType:
 215  		sig, ok := types.SafeUnderlying(argType).(*types.Signature)
 216  		if !ok {
 217  			return
 218  		}
 219  		for i, p := range pt.ParamList {
 220  			if sig.Params() != nil && i < int32(sig.Params().Len()) {
 221  				matchTypeParam(p.Type, sig.Params().At(int32(i)).Type(), tparams, typeMap)
 222  			}
 223  		}
 224  		for i, r := range pt.ResultList {
 225  			if sig.Results() != nil && i < int32(sig.Results().Len()) {
 226  				matchTypeParam(r.Type, sig.Results().At(int32(i)).Type(), tparams, typeMap)
 227  			}
 228  		}
 229  	}
 230  }
 231  
 232  func inferFromConstraints(tparams []*syntax.Field, typeMap map[string]types.Type) {
 233  	for _, tp := range tparams {
 234  		if tp.Name == nil || tp.Type == nil {
 235  			continue
 236  		}
 237  		knownType, ok := typeMap[tp.Name.Value]
 238  		if !ok {
 239  			continue
 240  		}
 241  		extractConstraintTypeParams(tp.Type, knownType, tparams, typeMap)
 242  	}
 243  }
 244  
 245  func extractConstraintTypeParams(constraint syntax.Expr, concreteType types.Type, tparams []*syntax.Field, typeMap map[string]types.Type) {
 246  	if constraint == nil || concreteType == nil {
 247  		return
 248  	}
 249  	switch c := constraint.(type) {
 250  	case *syntax.Operation:
 251  		if c.Op == token.Tilde && c.Y == nil {
 252  			extractConstraintTypeParams(c.X, concreteType, tparams, typeMap)
 253  		}
 254  	case *syntax.SliceType:
 255  		if sl, ok := types.SafeUnderlying(concreteType).(*types.Slice); ok {
 256  			if name, ok2 := c.Elem.(*syntax.Name); ok2 {
 257  				for _, tp := range tparams {
 258  					if tp.Name != nil && tp.Name.Value == name.Value {
 259  						if _, already := typeMap[name.Value]; !already {
 260  							typeMap[name.Value] = sl.Elem()
 261  						}
 262  					}
 263  				}
 264  			}
 265  		}
 266  	}
 267  }
 268  
 269  func inferTypeArgsFromFunc(fd *syntax.FuncDecl, args []SSAValue) (ts []types.Type) {
 270  	typeMap := map[string]types.Type{}
 271  	if fd.Type != nil {
 272  		for i, param := range fd.Type.ParamList {
 273  			if i >= len(args) || args[i] == nil || args[i].SSAType() == nil {
 274  				continue
 275  			}
 276  			matchTypeParam(param.Type, args[i].SSAType(), fd.TParamList, typeMap)
 277  		}
 278  	}
 279  	inferFromConstraints(fd.TParamList, typeMap)
 280  	var result []types.Type
 281  	for _, tp := range fd.TParamList {
 282  		if tp.Name != nil {
 283  			result = append(result, typeMap[tp.Name.Value])
 284  		}
 285  	}
 286  	return result
 287  }
 288  
 289  func (fb *ssaFuncBuilder) monomorphizeAndCall(fd *syntax.FuncDecl, baseName string, typeArgs []types.Type, args []SSAValue, targetPkg *SSAPackage, srcPkgPath string, hasDots bool) (s SSAValue) {
 290  	for _, ta := range typeArgs {
 291  		if ta == nil {
 292  			return nil
 293  		}
 294  	}
 295  	mangledName := mangleGenericName(baseName, typeArgs)
 296  	if fn, ok := targetPkg.Members[mangledName].(*SSAFunction); ok {
 297  		return fb.emitCallToSSAFunc(fn, args, hasDots)
 298  	}
 299  	typeSubst := map[string]types.Type{}
 300  	for i, tp := range fd.TParamList {
 301  		if tp.Name != nil && i < len(typeArgs) && typeArgs[i] != nil {
 302  			typeSubst[tp.Name.Value] = typeArgs[i]
 303  		}
 304  	}
 305  	var srcScope *types.Scope
 306  	if genericPkgScopes != nil {
 307  		if sc, ok := genericPkgScopes[srcPkgPath]; ok {
 308  			srcScope = sc
 309  		}
 310  	}
 311  	sig := resolveSignatureWithSubst(fd.Type, typeSubst, srcScope)
 312  	if sig == nil {
 313  		return nil
 314  	}
 315  	fn := &SSAFunction{
 316  		name:      mangledName,
 317  		Signature: sig,
 318  		Pkg:       targetPkg,
 319  		Prog:      fb.fn.Prog,
 320  	}
 321  	targetPkg.Members[mangledName] = fn
 322  	subFb := newSSAFuncBuilder(fn, fb.info)
 323  	subFb.typeSubst = typeSubst
 324  	subFb.srcScope = srcScope
 325  	subFb.buildBody(fd)
 326  	return fb.emitCallToSSAFunc(fn, args, hasDots)
 327  }
 328  
 329  func (fb *ssaFuncBuilder) emitCallToSSAFunc(fn *SSAFunction, args []SSAValue, hasDots bool) (s SSAValue) {
 330  	if fn.Signature != nil {
 331  		if fn.Signature.Variadic() && !hasDots {
 332  			args = fb.wrapVariadicArgs(args, fn.Signature)
 333  		}
 334  		args = fb.coerceArgsToInterface(args, fn.Signature)
 335  	}
 336  	var retType types.Type
 337  	if fn.Signature != nil && fn.Signature.Results() != nil {
 338  		if fn.Signature.Results().Len() == 1 {
 339  			retType = fn.Signature.Results().At(0).Type()
 340  		} else if fn.Signature.Results().Len() > 1 {
 341  			retType = fn.Signature.Results()
 342  		}
 343  	}
 344  	call := &SSACall{Call: SSACallCommon{Value: fn, Args: args}}
 345  	call.typ = retType
 346  	call.name = fb.nextName()
 347  	fb.emit(call)
 348  	if retType == nil {
 349  		return nil
 350  	}
 351  	return call
 352  }
 353  
 354  func (fb *ssaFuncBuilder) tryGenericLocalCall(funcName string, argList []syntax.Expr, hasDots bool) (s SSAValue) {
 355  	pkgPath := fb.fn.Pkg.Pkg.Path()
 356  	key := pkgPath | "." | funcName
 357  	fd, ok := genericFuncDecls[key]
 358  	if !ok && fb.srcScope != nil {
 359  		for gPkg, gScope := range genericPkgScopes {
 360  			if gScope == fb.srcScope {
 361  				key = gPkg | "." | funcName
 362  				fd, ok = genericFuncDecls[key]
 363  				if ok {
 364  					pkgPath = gPkg
 365  					break
 366  				}
 367  			}
 368  		}
 369  	}
 370  	if !ok {
 371  		return nil
 372  	}
 373  	args := fb.buildArgs(argList)
 374  	typeArgs := inferTypeArgsFromFunc(fd, args)
 375  	return fb.monomorphizeAndCall(fd, funcName, typeArgs, args, fb.fn.Pkg, pkgPath, hasDots)
 376  }
 377  
 378  func (fb *ssaFuncBuilder) tryGenericPkgCall(pn *types.PkgName, funcName string, argList []syntax.Expr, hasDots bool) (s SSAValue) {
 379  	imported := pn.Imported()
 380  	if imported == nil {
 381  		return nil
 382  	}
 383  	key := imported.Path() | "." | funcName
 384  	fd, ok := genericFuncDecls[key]
 385  	if !ok {
 386  		return nil
 387  	}
 388  	args := fb.buildArgs(argList)
 389  	typeArgs := inferTypeArgsFromFunc(fd, args)
 390  	return fb.monomorphizeAndCall(fd, funcName, typeArgs, args, fb.fn.Pkg, imported.Path(), hasDots)
 391  }
 392  
 393  func (fb *ssaFuncBuilder) tryGenericCallFromIndex(ie *syntax.IndexExpr, argList []syntax.Expr, hasDots bool) (s SSAValue) {
 394  	switch x := ie.X.(type) {
 395  	case *syntax.Name:
 396  		pkgPath := fb.fn.Pkg.Pkg.Path()
 397  		key := pkgPath | "." | x.Value
 398  		fd, ok := genericFuncDecls[key]
 399  		if !ok && fb.srcScope != nil {
 400  			for gPkg, gScope := range genericPkgScopes {
 401  				if gScope == fb.srcScope {
 402  					key = gPkg | "." | x.Value
 403  					fd, ok = genericFuncDecls[key]
 404  					if ok {
 405  						pkgPath = gPkg
 406  						break
 407  					}
 408  				}
 409  			}
 410  		}
 411  		if !ok {
 412  			return nil
 413  		}
 414  		typeArgs := fb.resolveExplicitTypeArgs(ie.Index, fd)
 415  		args := fb.buildArgs(argList)
 416  		return fb.monomorphizeAndCall(fd, x.Value, typeArgs, args, fb.fn.Pkg, pkgPath, hasDots)
 417  	case *syntax.SelectorExpr:
 418  		if pkgName, ok2 := x.X.(*syntax.Name); ok2 {
 419  			obj := fb.lookupObject(pkgName.Value)
 420  			if pn, ok3 := obj.(*types.PkgName); ok3 && pn.Imported() != nil {
 421  				key := pn.Imported().Path() | "." | x.Sel.Value
 422  				fd, ok4 := genericFuncDecls[key]
 423  				if !ok4 {
 424  					return nil
 425  				}
 426  				typeArgs := fb.resolveExplicitTypeArgs(ie.Index, fd)
 427  				args := fb.buildArgs(argList)
 428  				return fb.monomorphizeAndCall(fd, x.Sel.Value, typeArgs, args, fb.fn.Pkg, pn.Imported().Path(), hasDots)
 429  			}
 430  		}
 431  	}
 432  	return nil
 433  }
 434  
 435  func (fb *ssaFuncBuilder) resolveExplicitTypeArgs(indexExpr syntax.Expr, fd *syntax.FuncDecl) (ts []types.Type) {
 436  	var typeExprs []syntax.Expr
 437  	if le, ok := indexExpr.(*syntax.ListExpr); ok {
 438  		typeExprs = le.ElemList
 439  	} else {
 440  		typeExprs = []syntax.Expr{indexExpr}
 441  	}
 442  	var result []types.Type
 443  	for _, te := range typeExprs {
 444  		t := fb.resolveTypeAST(te)
 445  		result = append(result, t)
 446  	}
 447  	if len(result) < len(fd.TParamList) {
 448  		typeMap := map[string]types.Type{}
 449  		for i, tp := range fd.TParamList {
 450  			if tp.Name != nil && i < len(result) && result[i] != nil {
 451  				typeMap[tp.Name.Value] = result[i]
 452  			}
 453  		}
 454  		inferFromConstraints(fd.TParamList, typeMap)
 455  		result = nil
 456  		for _, tp := range fd.TParamList {
 457  			if tp.Name != nil {
 458  				result = append(result, typeMap[tp.Name.Value])
 459  			}
 460  		}
 461  	}
 462  	return result
 463  }
 464  
 465  func exprTypeName(e syntax.Expr) (s string) {
 466  	switch e := e.(type) {
 467  	case *syntax.Name:
 468  		return "Name(" | e.Value | ")"
 469  	case *syntax.BasicLit:
 470  		return "BasicLit"
 471  	case *syntax.SelectorExpr:
 472  		if x, ok := e.X.(*syntax.Name); ok {
 473  			return "Selector(" | x.Value | "." | e.Sel.Value | ")"
 474  		}
 475  		return "Selector"
 476  	case *syntax.CallExpr:
 477  		return "Call"
 478  	case *syntax.Operation:
 479  		return "Op"
 480  	}
 481  	return "unknown"
 482  }
 483