fix StrLiteral and BytesLiteral to support constant (#11)

[bytom/equity.git] / compiler / compile.go

diff --git a/compiler/compile.go b/compiler/compile.go
index b9199f2..76913a2 100644 (file)
--- a/compiler/compile.go
+++ b/compiler/compile.go
@@ -5,7 +5,6 @@ import (
        "fmt"
        "io"
        "io/ioutil"
-       "strings"
 
        chainjson "github.com/bytom/encoding/json"
        "github.com/bytom/errors"
@@ -321,6 +320,10 @@ func compileClause(b *builder, contractStk stack, contract *Contract, env *envir
        // a count of the number of times each variable is referenced
        counts := make(map[string]int)
        for _, s := range clause.statements {
+               if stmt, ok := s.(*defineStatement); ok && stmt.expr == nil {
+                       continue
+               }
+
                s.countVarRefs(counts)
                if stmt, ok := s.(*ifStatement); ok {
                        for _, trueStmt := range stmt.body.trueBody {
@@ -449,19 +452,74 @@ func compileStatement(b *builder, stk stack, contract *Contract, env *environ, c
                b.addJumpTarget(stk, "endif_"+strSequence)
 
        case *defineStatement:
+               // add environ for define variable
+               if err = env.add(stmt.variable.Name, stmt.variable.Type, roleClauseVariable); err != nil {
+                       return stk, err
+               }
+
+               // check whether the variable is used or not
+               if counts[stmt.variable.Name] == 0 {
+                       return stk, fmt.Errorf("the defined variable \"%s\" is unused in clause \"%s\"", stmt.variable.Name, clause.Name)
+               }
+
+               if stmt.expr != nil {
+                       // variable
+                       stk, err = compileExpr(b, stk, contract, clause, env, counts, stmt.expr)
+                       if err != nil {
+                               return stk, errors.Wrapf(err, "in define statement in clause \"%s\"", clause.Name)
+                       }
+
+                       // modify stack name
+                       stk.str = stmt.variable.Name
+               }
+
+       case *assignStatement:
+               // find variable from environ with roleClauseVariable
+               if entry := env.lookup(string(stmt.variable.Name)); entry != nil {
+                       if entry.r != roleClauseVariable {
+                               return stk, fmt.Errorf("the type of variable is not roleClauseVariable in assign statement in clause \"%s\"", clause.Name)
+                       }
+                       stmt.variable.Type = entry.t
+               } else {
+                       return stk, fmt.Errorf("the variable \"%s\" is not defined before the assign statement in clause \"%s\"", stmt.variable.Name, clause.Name)
+               }
+
+               // temporary store the counts of defined variable
+               varCount := counts[stmt.variable.Name]
+
+               // calculate the counts of variable for assign statement
+               tmpCounts := make(map[string]int)
+               stmt.countVarRefs(tmpCounts)
+
+               // modify the map counts of defined variable to 1 and minus the number of defined variable
+               // when the assign expression contains the defined variable
+               if tmpCounts[stmt.variable.Name] > 0 {
+                       counts[stmt.variable.Name] = 1
+                       varCount -= tmpCounts[stmt.variable.Name]
+               } else {
+                       depth := stk.find(stmt.variable.Name)
+                       switch depth {
+                       case 0:
+                               break
+                       case 1:
+                               stk = b.addSwap(stk)
+                       default:
+                               stk = b.addRoll(stk, depth)
+                       }
+                       stk = b.addDrop(stk)
+               }
+
                // variable
                stk, err = compileExpr(b, stk, contract, clause, env, counts, stmt.expr)
                if err != nil {
                        return stk, errors.Wrapf(err, "in define statement in clause \"%s\"", clause.Name)
                }
 
-               // modify stack name
-               stk.str = stmt.varName.Name
+               // restore the defined variable counts
+               counts[stmt.variable.Name] = varCount
 
-               // add environ for define variable
-               if err = env.add(stmt.varName.Name, stmt.varName.Type, roleClauseVariable); err != nil {
-                       return stk, err
-               }
+               // modify stack name
+               stk.str = stmt.variable.Name
 
        case *verifyStatement:
                stk, err = compileExpr(b, stk, contract, clause, env, counts, stmt.expr)
@@ -493,16 +551,21 @@ func compileStatement(b *builder, stk stack, contract *Contract, env *environ, c
                        stk = b.addAmount(stk, contract.Value.Amount)
                        stk = b.addAsset(stk, contract.Value.Asset)
                } else {
+                       // calculate the counts of variable for lockStatement
+                       lockCounts := make(map[string]int)
+                       stmt.countVarRefs(lockCounts)
+
                        // amount
-                       if stmt.lockedAmount.String() == contract.Value.Amount {
+                       switch {
+                       case stmt.lockedAmount.String() == contract.Value.Amount:
                                stk = b.addAmount(stk, contract.Value.Amount)
-                       } else if strings.Contains(stmt.lockedAmount.String(), contract.Value.Amount) {
+                       case stmt.lockedAmount.String() != contract.Value.Amount && lockCounts[contract.Value.Amount] > 0:
                                stk = b.addAmount(stk, contract.Value.Amount)
                                stk, err = compileExpr(b, stk, contract, clause, env, counts, stmt.lockedAmount)
                                if err != nil {
                                        return stk, errors.Wrapf(err, "in lock statement in clause \"%s\"", clause.Name)
                                }
-                       } else {
+                       default:
                                stk, err = compileExpr(b, stk, contract, clause, env, counts, stmt.lockedAmount)
                                if err != nil {
                                        return stk, errors.Wrapf(err, "in lock statement in clause \"%s\"", clause.Name)
@@ -510,15 +573,16 @@ func compileStatement(b *builder, stk stack, contract *Contract, env *environ, c
                        }
 
                        // asset
-                       if stmt.lockedAsset.String() == contract.Value.Asset {
+                       switch {
+                       case stmt.lockedAsset.String() == contract.Value.Asset:
                                stk = b.addAsset(stk, contract.Value.Asset)
-                       } else if strings.Contains(stmt.lockedAsset.String(), contract.Value.Asset) {
+                       case stmt.lockedAsset.String() != contract.Value.Asset && lockCounts[contract.Value.Asset] > 0:
                                stk = b.addAsset(stk, contract.Value.Asset)
                                stk, err = compileExpr(b, stk, contract, clause, env, counts, stmt.lockedAsset)
                                if err != nil {
                                        return stk, errors.Wrapf(err, "in lock statement in clause \"%s\"", clause.Name)
                                }
-                       } else {
+                       default:
                                stk, err = compileExpr(b, stk, contract, clause, env, counts, stmt.lockedAsset)
                                if err != nil {
                                        return stk, errors.Wrapf(err, "in lock statement in clause \"%s\"", clause.Name)