package jasmin

  1. Overview
  2. Docs
package jasmin
Legend:
Page
Library
Module
Module type
Parameter
Class
Class type
Source

Source file sct_checker_forward.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
open Utils
open Annotations
open Prog
open Constraints

module CT = Ct_checker_forward

module S = Syntax

(* ----------------------------------------------------------- *)
let pp_var fmt x = Printer.pp_var ~debug:false fmt x

let pp_var_i fmt x = pp_var fmt (L.unloc x)

let pp_expr fmt e = Printer.pp_expr ~debug:false fmt e

let pp_lval fmt x = Printer.pp_lval ~debug:false fmt x

let pp_vset fmt xs =
  Format.fprintf fmt "{@[ %a @]}"  (pp_list ",@ " pp_var) (Sv.elements xs)


(* ----------------------------------------------------------- *)
let ssecret = "secret"
let spublic = "public"
let stransient = "transient"
let smsf = "msf"

(* ----------------------------------------------------------- *)
(* Info provided by the user                                   *)
type ulevel =
  | Secret
  | Transient
  | Public
  | Msf

(* -------------------------------------------------------------- *)
(* Special operators to deal with msf                             *)
type special_op =
  | Init_msf
  | Update_msf
  | Mov_msf
  | Protect
  | Other

let is_special o =
  match o with
  | Sopn.Opseudo_op (Pseudo_operator.Ospill _) -> assert false
  | Sopn.Opseudo_op _ -> Other
  | Oasm _ -> Other
  | Oslh o ->
    match o with
    | SLHinit   -> Init_msf
    | SLHupdate -> Update_msf
    | SLHmove   -> Mov_msf
    | SLHprotect _ | SLHprotect_ptr _ -> Protect
    | SLHprotect_ptr_fail _ -> assert false

(* -----------------------------------------------------------*)
(* Types generated by the system                              *)

(* Attempt of convention: n stands for normal (non-speculative),
   s for speculative *)

type vty =
  | Direct of VlPairs.t
  | Indirect of VlPairs.t * VlPairs.t

let pp_vlpair ~top fmt (n, s) =
  if Lvl.is_public n && Lvl.is_public s then
    Format.fprintf fmt "%s%s" (if top then "#" else "") spublic
  else if Lvl.is_public n && Lvl.is_secret s then Format.fprintf fmt "%s%s" (if top then "#" else "") stransient
  else if Lvl.is_secret n && Lvl.is_secret s then Format.fprintf fmt "%s%s" (if top then "#" else "")ssecret
  else
    Format.fprintf fmt "%s{ n = %a, s = %a}" (if top then "#poly = " else "") Lvl.pp n Lvl.pp s

let pp_vty fmt = function
  | Direct vlp -> pp_vlpair ~top:true fmt vlp
  | Indirect(p_vlp, v_vlp) ->
      Format.fprintf fmt "#[ptr = %a, val = %a]" (pp_vlpair ~top:false) p_vlp (pp_vlpair ~top:false) v_vlp

type vfty =
  | IsMsf
  | IsNormal of vty

let pp_vfty fmt = function
  | IsMsf -> Format.fprintf fmt "#%s" smsf
  | IsNormal ty -> pp_vty fmt ty

(* Either a function does not modify the MSF, or it does at a certain location,
   and if that location is a function call we add the trace of calls until the
   offending instruction. *)
type modmsf =
  | Modified of L.i_loc * (L.i_loc * funname) list
  | NotModified

let is_Modified m =
  match m with
  | Modified _ -> true
  | NotModified -> false

type ty_fun = {
    modmsf               : modmsf;
    tyin                 : vfty list;
    tyout                : vfty list;
    constraints          : C.constraints;
    resulting_corruption : VlPairs.t (* resulting memory corruption after function call *)
  }

type ('info,'asm) fenv = {
    env_ty     : ty_fun Hf.t;
    env_def    : ('info,'asm) func list;
  }

module FEnv = struct
  let get_fun_def fenv fn = List.find (fun f -> F.equal f.f_name fn) fenv.env_def

  let get_fty fenv fn =
    try Hf.find fenv.env_ty fn with Not_found -> assert false

end

let pp_modmsf fmt modmsf =
  let s =
    match modmsf with
    | Modified _ -> "modmsf"
    | NotModified -> "nomodmsf"
  in
  Format.fprintf fmt "%s" s

let pp_funty fmt (fname, tyfun) =
  Format.fprintf fmt
    "@[<v>%a %s : @[%a@] ->@ @[%a@]@ \
      output corruption: %a@.\
      @ constraints:@ @[%a@]@]@."
    pp_modmsf tyfun.modmsf
    fname
    (pp_list " *@ " pp_vfty) tyfun.tyin
    (pp_list " *@ " pp_vfty) tyfun.tyout
    pp_vty (Direct (tyfun.resulting_corruption))
    C.pp tyfun.constraints


(* --------------------------------------------------------------- *)
(* Checks if the status of msf is modified by the instructions     *)

let is_inline i =
  match Annot.ensure_uniq1 "inline" Annot.none i.i_annot with
  | Some _ -> true
  | None   -> false

let rec modmsf_i fenv i =
  let modified_here = Modified(i.i_loc, []) in
  match i.i_desc with
  | Csyscall _ | Cwhile _ -> modified_here
  | Cif(_, c0, c1) ->
    if is_inline i
    then
      let r = modmsf_c fenv c0 in
      if is_Modified r then r else modmsf_c fenv c1
    else modified_here
  | Cassgn _ -> NotModified
  | Copn (_, _, o, _) ->
    begin match is_special o with
    | Init_msf -> modified_here (* LFENCE modifies msf *)
    | Update_msf -> modified_here (* not sure it is needed *)
    | Mov_msf | Protect | Other -> NotModified
    end
  | Cfor(_, _, c) -> modmsf_c fenv c
  | Ccall (_, f, _) ->
    match (FEnv.get_fty fenv f).modmsf with
    | Modified (l, tr) -> Modified(i.i_loc, (l, f) :: tr)
    | NotModified -> NotModified

and modmsf_c fenv c =
  List.map (modmsf_i fenv) c
  |> List.find_opt is_Modified
  |> Option.default NotModified

let error ~loc =
  hierror ~loc:(Lone loc) ~kind:"speculative constant type checker"

let warn ~loc = warning SCTchecker (L.i_loc0 loc)

(* --------------------------------------------------------- *)
(* Inference of the variables that need to contain msf       *)
(* This code does not need to be trusted it is simply        *)
(* used as an oracle                                         *)

let is_register ~direct x =
  is_reg_kind x.v_kind && (not direct || is_reg_direct_kind x.v_kind)

let ensure_register ~direct x =
  if not (is_register ~direct (L.unloc x)) then
    error ~loc:(L.loc x)
      "variable %a must be a reg%s" pp_var_i x
      (if direct then "" else " (ptr)")

let reg_lval ~direct loc x =
  match x with
  | Laset (_, _, _, x, _)
  | Lvar x -> ensure_register ~direct x; x
  | _ ->
      error ~loc "L-value %a must be a reg%s" pp_lval x
        (if direct then "" else " (ptr)")

let reg_lval_opt ~direct loc =
  function
  | Lnone _ -> None
  | x -> Some (reg_lval ~direct loc x)

let reg_expr_opt ~direct = function
  | Pget (_, _, _, x, _) | Pvar x ->
      if is_gkvar x && is_register ~direct (L.unloc x.gv)
      then Some x.gv
      else None
  | _ -> None

let reg_expr ~direct loc e =
  match reg_expr_opt ~direct e with
  | Some x -> x
  | None -> error ~loc "expression %a must be a reg%s" pp_expr e
        (if direct then "" else " (ptr)")


let rec infer_msf_i ~withcheck fenv (tbl:(L.i_loc, Sv.t) Hashtbl.t) i ms =
  let loc = i.i_loc.L.base_loc in

  let check_x ms x =
    if Sv.mem (L.unloc x) ms then
      error ~loc:(L.loc x)
        "instruction assigns %a, which is required to be a msf"
        pp_var_i x in

  let check ms x =
    match x with
    | Lvar x -> check_x ms x
    | _ -> () in

  let checks ms xs = List.iter (check ms) xs in

  let pp_modmsf_trace fmt tr =
    let pp_item fmt (l, fn) =
      Format.fprintf fmt
        "@[<h>the function %s destroys MSFs at %a@]"
        fn.fn_name
        L.pp_iloc l
    in
    Format.fprintf fmt "Trace:@;<0 2>@[<v>%a@]" (pp_list "" pp_item) tr
  in

  let check_call ~loc fn modmsf ms =
    match modmsf with
    | Modified(l, tr) ->
      if not (Sv.is_empty ms) && withcheck then
        error ~loc
          "@[<h>this function call destroys MSFs and %a are required.@]@;%a"
          pp_vset ms
          pp_modmsf_trace ((l, fn) :: tr)
    | NotModified -> ()
  in

  match i.i_desc with
  | Csyscall _ ->
      if not (Sv.is_empty ms) && withcheck then
        error ~loc "syscalls destroy msf variables, %a are required" pp_vset ms;
      (* withcheck => is_empty ms *)
      ms

  | Cif (_, c1, c2) ->
    let ms1 = infer_msf_c ~withcheck fenv tbl c1 ms in
    let ms2 = infer_msf_c ~withcheck fenv tbl c2 ms in
    Sv.union ms1 ms2

  | Cfor(x, _, c) ->
    check_x ms x;
    let rec loop ms =
      let ms' = infer_msf_c ~withcheck fenv tbl c ms in
      if Sv.subset ms' ms then (Hashtbl.add tbl i.i_loc ms; ms)
      else loop (Sv.union ms' ms) in
    loop ms

  | Cwhile (_, c1, _, _, c2) ->
    (* c1; while e do c2; c1 *)
    let rec loop ms =
      let ms1 = infer_msf_c ~withcheck fenv tbl c1 ms in
      let ms2 = infer_msf_c ~withcheck fenv tbl c2 ms1 in
      if Sv.subset ms2 ms then (Hashtbl.add tbl i.i_loc ms1; ms1)
      else loop (Sv.union ms2 ms) in
    loop ms

  | Cassgn(Lvar x, tag, _, e) when Sv.mem (L.unloc x) ms ->
      (* We need to allow assignments to MSF if the compiler introduces them,
         to be able to use the checker after inlining. *)
      let gets_removed =
        match tag with
        | AT_none | AT_keep -> false
        | AT_rename | AT_inline | AT_phinode -> true
      in
      begin match reg_expr_opt ~direct:true e with
      | Some x' when gets_removed ->
          Sv.add (L.unloc x') (Sv.remove (L.unloc x) ms)
      | _ -> error ~loc "assignment to MSF variable %a not allowed" pp_var_i x
      end

  | Cassgn _ ->
    ms

  | Ccall(xs, f, es) ->
    let fty = FEnv.get_fty fenv f in
    let ms =
      let doout ms vfty x =
        match vfty with
        | IsMsf -> let x = reg_lval ~direct:true loc x in Sv.remove (L.unloc x) ms
        | _ -> ms in
      let ms = List.fold_left2 doout ms fty.tyout xs in
      check_call ~loc f fty.modmsf ms;
      ms
    in
    let doin ms vfty e =
      match vfty with
      | IsMsf -> let x = reg_expr ~direct:true loc e in Sv.add (L.unloc x) ms
      | _ -> ms in
    List.fold_left2 doin ms fty.tyin es

  | Copn (xs, _, o, es) ->
    match is_special o, xs, es with
    | Init_msf, [x], _ ->
      let x = Option.map_default (fun x -> Sv.singleton (L.unloc x)) Sv.empty (reg_lval_opt ~direct:true loc x) in
      if not (Sv.subset ms x) then
        error ~loc "only %a is ensured to be msf after init_msf, %a are required" pp_vset x pp_vset ms;
      Sv.empty

    | Init_msf, _, _ -> assert false

    | Update_msf, [xms], [_e; ms'] ->
      let xms = reg_lval ~direct:true loc xms and ms' = reg_expr ~direct:true loc ms' in
      if not (Sv.subset ms (Sv.singleton (L.unloc xms))) then
        error ~loc "only %a is ensured to be msf after update_msf, %a are required" pp_var_i xms pp_vset ms;
      Sv.singleton (L.unloc ms')

    | Update_msf, _, _ -> assert false

    | Mov_msf, [x], [x'] ->
      let x = reg_lval ~direct:true loc x and x' = reg_expr ~direct:true loc x' in
      Sv.add (L.unloc x') (Sv.remove (L.unloc x) ms)

    | Mov_msf, _, _ -> assert false

    | Protect, xs, [_; ms'] ->
      checks ms xs;
      let ms' = reg_expr ~direct:true loc ms' in
      Sv.add (L.unloc ms') ms

    | Protect, _, _ -> assert false

    | Other, xs, _ -> checks ms xs; ms

and infer_msf_c ~withcheck fenv tbl c ms =
  List.fold_right (infer_msf_i ~withcheck fenv tbl) c ms


(* --------------------------------------------------------- *)
(* Typing environment                                        *)

module Env : sig
  type env  (* constraints holder *)
  type venv (* type variables association *)

  val init : unit -> env

  val empty : env -> venv
  val constraints : env -> C.constraints
  val add_var : env -> venv -> var -> vty -> venv
  val public    : env -> Lvl.t
  val secret    : env -> Lvl.t

  val public2    : env -> VlPairs.t
  val transient  : env -> VlPairs.t
  val secret2    : env -> VlPairs.t

  val dpublic    : env -> vty
  val dsecret    : env -> vty

  val get   : venv -> var -> vty
  val get_i : venv -> var_i -> vty
  val gget  : venv -> int ggvar -> vty

  val fresh  : ?name:string -> env -> Lvl.t
  val fresh2 : ?name:string -> env -> VlPairs.t

  val init_ty : env -> venv -> var -> vty -> venv
  val set_ty : env -> venv -> var_i -> vty -> venv
  val set_init_msf : env -> venv -> var_i option -> venv

  val max : env -> venv -> venv -> venv

  val get_msf_oracle : env -> (L.i_loc, Sv.t) Hashtbl.t
  val msf_oracle : env -> L.i_loc -> Sv.t

  val freshen : ?min:Constraints.VlPairs.t -> env -> venv -> venv
  val ensure_le : L.t -> venv -> venv -> unit
  val clone_for_call : env -> ty_fun -> vfty list * vfty list * VlPairs.t
          (* output type, input type, output corruption *)

  val corruption : env -> venv -> VlPairs.t -> venv
  val corruption_speculative : env -> venv -> VlPairs.t -> venv

  val get_resulting_corruption : venv -> VlPairs.t

  val venv_forget : Sv.t -> venv -> venv
  (** Restrict knowledge to the given set of variables. *)

end = struct

  type env = {
      constraints : C.constraints;
      msf_oracle  : (L.i_loc, Sv.t) Hashtbl.t;
    }

  type venv = {
      vtype : vty Mv.t;
      vars : Sv.t;
      resulting_corruption : VlPairs.t;
      public2 : VlPairs.t (* needed in 'get'. This is dirty.
        Ideally, env and venv should be merged, but this requires
        some work and restructuring of the module signature. *)
    }

  let init () =
    { constraints = C.init ();
      msf_oracle  = Hashtbl.create 97; }

  let constraints env = env.constraints

  let public    env = C.public    env.constraints
  let secret    env = C.secret    env.constraints

  let public2   env = (public env, public env)
  let transient env = (public env, secret env)
  let secret2   env = (secret env, secret env)

  let dpublic   env = Direct (public2 env)
  let dsecret   env = Direct (secret2 env)

  let fresh ?name env = C.fresh ?name env.constraints

  let fresh2 ?name env = (C.fresh ?name env.constraints,
                          C.fresh ?name env.constraints)

  let empty env =
    { vtype = Mv.empty;
      vars = Sv.empty;
      resulting_corruption = fresh2 env;
      public2 = public2 env;
    }

  let get venv x =
    try match x.v_kind with
    | Global -> Direct venv.public2
    | _ -> Mv.find x venv.vtype
    with Not_found -> assert false

  let get_i venv x = get venv (L.unloc x)
  let gget venv x = get_i venv x.gv

  let add_le_var ty1 ty2 =
    match ty1, ty2 with
    | Direct ty1, Direct ty2 -> VlPairs.add_le ty1 ty2
    | Indirect (l1, x1), Indirect (l2, x2) -> VlPairs.add_le l1 l2; VlPairs.add_le x1 x2
    | _ -> assert false

  let init_ty env venv x xty =
      (* warning Always (L.i_loc0 x.v_dloc) "initialising variable %a" pp_var x; *)
    let nxty =
      let public2 = public2 env in
      match x.v_kind, xty with
      | Const, Direct le -> VlPairs.add_le le public2; xty
      | Inline, Direct le -> VlPairs.add_le le public2; xty

      | Global, Direct _
      | Stack Direct, Direct _
      | Stack (Pointer _), Indirect _
      | Reg (_, Direct), Direct _
      | Reg (_, Pointer _), Indirect _ -> xty

      | _, _ -> assert false
    in
    { venv with
        vtype = Mv.add x nxty venv.vtype;
        vars = Sv.add x venv.vars
    }

  let add_var env venv x vty =
    assert (not (Mv.mem x venv.vtype));
    try init_ty env venv x vty
    with Lvl.Unsat _unsat ->
      error ~loc:(x.v_dloc)
        "invalid security annotations for %a, this leads to invalid constraints"
           pp_var x

  let set_ty env venv x nxty =
    init_ty env venv (L.unloc x) nxty

  (* Initialises type for InitMsf operations. Acts as a Fence operator *)
  let set_init_msf env venv ms =
    let venv =
      match ms with
      | Some ms -> set_ty env venv ms (dpublic env)
      | None -> venv in
    let operate_fence _ = function
        | Direct le -> Direct (VlPairs.normalise le)
        | Indirect(lp, le) -> Indirect(VlPairs.normalise lp, VlPairs.normalise le)
    in
    { venv with vtype = Mv.mapi operate_fence venv.vtype }


  let max env venv1 venv2 =
    let merge1 l1 l2 =
      if Lvl.equal l1 l2 then l1
      else
        let l = fresh env in
        Lvl.add_le l1 l; Lvl.add_le l2 l;
        l in

    let merge (n1, s1) (n2, s2) =
      (merge1 n1 n2, merge1 s1 s2) in

    let merge_var _ oty1 oty2 =
      match oget oty1, oget oty2 with
      | Direct le1, Direct le2 -> Some (Direct (merge le1 le2))
      | Indirect(lp1, le1), Indirect(lp2, le2) -> Some (Indirect (merge lp1 lp2, merge le1 le2))
      | _ -> assert false
    in
    { venv1 with
      vtype = Mv.merge merge_var venv1.vtype venv2.vtype;
      resulting_corruption = merge venv1.resulting_corruption venv2.resulting_corruption
    }

  let get_msf_oracle env = env.msf_oracle

  let msf_oracle env loc =
    try Hashtbl.find env.msf_oracle loc with Not_found -> assert false

  (* freshen all variables in environment env, venv, with
     possibly a minimum (typically memory corruption) *)
  let freshen ?min env venv =
    let fresh ~in_memory le =
      let l = fresh2 env in
      if in_memory && min != None then VlPairs.add_le (oget min) l;
      VlPairs.add_le le l;
      l
    in
    { venv with vtype = Sv.fold (fun x vtype ->
        let in_memory = match x.v_kind with
          | Wsize.Global (* likely unused as global variables are not in venv.vars *)
          | Stack _ -> true
          | Const | Inline | Reg _ -> false
        in
        let ty =
          match Mv.find x vtype with
          | Direct le -> Direct (fresh ~in_memory le)
          | Indirect(lp, le) ->
             Indirect(fresh ~in_memory lp, fresh ~in_memory:true le) (* the pointed values are in memory *)
        in
        Mv.add x ty vtype) venv.vars venv.vtype }

  let ensure_le loc venv1 venv2 =
    let add_le_silent _ oty1 oty2 = add_le_var (oget oty1) (oget oty2); None in
    try ignore (Mv.merge add_le_silent venv1.vtype venv2.vtype)
    with Lvl.Unsat _unsat ->
      error ~loc "constraints caused by the loop cannot be satisfied"

  let clone_for_call (env:env) (tyfun:ty_fun) =
    let subst1 = C.clone tyfun.constraints env.constraints in
    let subst (n, s) = (subst1 n, subst1 s) in
    let subst_ty = function
      | IsMsf -> IsMsf
      | IsNormal ty -> let ty =
            match ty with
            | Direct le -> Direct (subst le)
            | Indirect(lp, le) -> Indirect(subst lp, subst le) in
          IsNormal ty in
    List.map subst_ty tyfun.tyout, List.map subst_ty tyfun.tyin,
    subst tyfun.resulting_corruption

  let corruption env venv ty =
    VlPairs.add_le ty venv.resulting_corruption; (* update corruption level *)
    freshen ~min:ty env venv

  let corruption_speculative env venv (_, s) = corruption env venv (public env, s)

  let get_resulting_corruption venv = venv.resulting_corruption

  let venv_forget lv x =
    { x with
      vtype = Mv.filter (fun k _ -> Sv.mem k lv) x.vtype;
      vars = Sv.inter x.vars lv }

end


(* --------------------------------------------------------- *)
let error_unsat loc (_ : Lvl.t list * Lvl.t * Lvl.t) pp e ety ety' =
  error ~loc
    "%a has type %a but should be at most %a"
    pp e pp_vty ety pp_vty ety'

let ssafe_test x aa ws i =
  let x = L.unloc x in
  match x.v_kind, x.v_ty, i with
  | Reg (_, Direct), _, _ -> true
  | _, Arr (ws1, len), Pconst v ->
      let len = Z.of_int (arr_size ws1 len) in
      let v = Z.of_int (access_offset aa ws (Z.to_int v)) in
      let v_max = Z.add v (Z.of_int (size_of_ws ws - 1)) in
      Z.(leq zero v && lt v_max len)
  | _ -> false

let content_ty = function
  | Direct le | Indirect (_, le) -> le

(* --------------------------------------------------------- *)
(* Type checking of expressions                              *)

let rec ty_expr env venv loc (e:expr) : vty =
  match e with
  | Pconst _ | Pbool _ | Parr_init _ -> Env.dpublic env

  | Pvar x -> Env.gget venv x

  | Pget (_, aa, ws, x, i) ->
      ensure_public_address env venv loc x.gv;
      ensure_public env venv loc i;
      let ty = Env.fresh2 env
      and xty = Env.gget venv x in
      if not (ssafe_test x.gv aa ws i) then VlPairs.add_le_speculative (Env.secret env) ty;
      VlPairs.add_le (content_ty xty) ty;
      Direct ty

    (* in the case of sub-arrays, no operation is performed, and there is now
       an alias on the values. Thus the type must be *equal* *)
  | Psub (_, _, _, x, i) ->
      ensure_public env venv loc i;
      Env.gget venv x

  | Pload (_, _, i) ->
      ensure_public env venv loc i;
      Env.dsecret env

  | Papp1(o, e)      ->
    let public = not (CT.is_ct_op1 o) in
    ty_exprs_max ~public env venv loc [e]
  | Papp2(o, e1, e2) ->
    let public = not (CT.is_ct_op2 o) in
    ty_exprs_max ~public env venv loc [e1; e2]
  | PappN(o, es)     ->
    let public = not (CT.is_ct_opN o) in
    ty_exprs_max ~public env venv loc es

  | Pif(_, e1, e2, e3) ->
      let ty1 = ty_expr env venv loc e1 in
      let ty2 = ty_expr env venv loc e2 in
      let ty3 = ty_expr env venv loc e3 in
      match ty1 with
      | Indirect _ -> assert false
      | Direct l1 ->
        let do_indirect lp2 le2 lp3 le3 =
          let lp = Env.fresh2 env in
          let le = Env.fresh2 env in
          (* The condition expression is also added to the constraints because it can be deduced from the result value*)
          VlPairs.add_le l1 lp; VlPairs.add_le lp2 lp; VlPairs.add_le lp3 lp;
          VlPairs.add_le l1 le; VlPairs.add_le le2 le; VlPairs.add_le le3 le;
          Indirect(lp, le) in
        match ty2, ty3 with
        | Direct l2, Direct l3 ->
          let le = Env.fresh2 env in
          VlPairs.add_le l1 le; VlPairs.add_le l2 le; VlPairs.add_le l3 le;
          Direct le
        | Indirect(lp2, le2), Indirect(lp3, le3) ->
          do_indirect lp2 le2 lp3 le3
        | Indirect(lp2, le2), Direct le3 ->
          do_indirect lp2 le2 (Env.public2 env) le3
        | Direct le2, Indirect (lp3, le3) ->
          do_indirect (Env.public2 env) le2 lp3 le3

and ensure_smaller env venv loc e l =
  let ety = ty_expr env venv loc e in
  match ety with
  | Direct le | Indirect (le, _) ->
    try VlPairs.add_le le l
    with Lvl.Unsat unsat -> error_unsat loc unsat pp_expr e ety (Direct l)

and ensure_public env venv loc e = ensure_smaller env venv loc e (Env.public2 env)

and ensure_public_address env venv loc x =
  let ety = Env.get_i venv x in
  match ety with
  | Direct _ -> () (* stack or reg arrays have public addresses by definition *)
  | Indirect (le, _) -> try VlPairs.add_le le (Env.public2 env)
      with Lvl.Unsat unsat -> error_unsat loc unsat pp_var_i x ety (Direct (Env.public2 env))

and ty_exprs_max ~(public:bool) env venv loc es : vty =
  let l = if public then Env.public2 env else Env.fresh2 env in
  List.iter (fun e -> ensure_smaller env venv loc e l) es;
  Direct l


(* ------------------------------------------------------------- *)
(* Compare expressions up to constant folding *)
let expr_equal a b =
  let fcp =
    let open Glob_options in
    match !target_arch with
    | X86_64 -> X86_decl.x86_fcp
    | ARM_M4 -> Arm_decl.arm_fcp 
    | RISCV  -> Riscv_decl.riscv_fcp
  in 
  let normalize e =
    e |> Conv.cexpr_of_expr |> Constant_prop.(const_prop_e fcp None empty_cpm) in
  Expr.eq_expr (normalize a) (normalize b)

(* ------------------------------------------------------------- *)
(* MSF status                                                    *)

module MSF : sig

  type t
  val toinit : t
  val exact1 : var_i -> t
  val add : var_i -> t -> t
  val update   : t -> var  -> t

  val enter_if : t -> expr -> t
  val max : t -> t -> t

  val check_msf : t -> var_i -> unit
  val check_msf_trans : t -> var_i -> expr -> unit
  val is_msf    : t -> var_i -> bool
  val is_msf_exact    : t -> var_i -> bool
  val check_msf_exact : t -> var_i -> unit
  val loop : Env.env -> L.i_loc -> t -> t
  val end_loop : L.t -> t -> t -> t

  val pp : Format.formatter -> t -> unit

  end = struct

  type t = Sv.t * expr option

  let toinit = (Sv.empty, None)
  let exact xs = (xs, None)
  let trans xs e = (xs, Some e)

  let exact1 x =
    ensure_register ~direct:true x;
    exact (Sv.singleton (L.unloc x))

  let add x (xs, oe) =
    ensure_register ~direct:true x;
    let loc = L.loc x in
    let x = L.unloc x in
    Stdlib.Option.iter (fun e ->
        if Sv.mem x (vars_e e) then
          error ~loc "%a cannot become an MSF as the current status depends on it (%a)" pp_var x pp_expr e
      ) oe;
    (Sv.add x xs, oe)

  let update (xs, oe) x =
    match oe with
    | Some e when Sv.mem x (vars_e e) -> toinit
    | _ -> (Sv.remove x xs, oe)

  let enter_if msf e =
    match msf  with
    | (_, Some _) -> toinit
    | (xs, None) -> (xs, Some e)

  let max (xs1, oe1) (xs2, oe2) =
    match oe1, oe2 with
    | None, None -> Sv.inter xs1 xs2, None
    | Some e1, Some e2 when expr_equal e1 e2 -> Sv.inter xs1 xs2, Some e1
    | _, _ -> toinit

  let check_msf (xs, _) ms =
    if not (Sv.mem (L.unloc ms) xs) then
      error ~loc:(L.loc ms)
        "the variable %a is not known to be a msf, only %a are"
        pp_var_i ms pp_vset xs

  let check_msf_trans ((_, ob) as msf) ms b =
    match ob with
    | None -> error ~loc:(L.loc ms) "MSF is not Trans"
    | Some b' ->
        if not (expr_equal b b') then
          error ~loc:(L.loc ms)
          "the expression %a need to be equal to@ %a"
            pp_expr b pp_expr  b';
        check_msf msf ms

  let is_msf (xs, _) ms = Sv.mem (L.unloc ms) xs

  let is_msf_exact (xs, ob) ms =
    match ob with
    | Some _ -> false
    | None -> Sv.mem (L.unloc ms) xs

  let check_msf_exact ((_, ob) as msf) ms =
    match ob with
    | Some b ->
      error ~loc:(L.loc ms) "MSF is Trans@ %a"  pp_expr b
    | None -> check_msf msf ms

  let pp fmt (xs, oe) =
    match oe with
    | Some e -> Format.fprintf fmt "Trans %a %a" pp_vset xs pp_expr e
    | None   -> Format.fprintf fmt "Exact %a" pp_vset xs

  let loop env loc ((xs, oe) as msf) =
    let xs' = Env.msf_oracle env loc in
    if Sv.subset xs' xs then (xs', oe)
    else
      error ~loc:(loc.L.base_loc)
        "current msf = %a, it should contain %a"
        pp msf pp_vset (Sv.diff xs' xs)

  let end_loop loc ((xsi, oei) as msfi) ((xso, oeo) as msfo)=
    if Sv.subset xsi xso then
      begin match oei, oeo with
      | None, None -> msfi
      | Some ei, Some eo when expr_equal ei eo -> msfi
      | _, _ ->
          if not (Sv.is_empty xsi) then
            error ~loc
            "msf is %a it should be be at least %a"
            pp msfo pp msfi;
          toinit
    end
    else
      error ~loc
        "msf is %a it should be be at least %a"
        pp msfo pp msfi

end


(* --------------------------------------------------------- *)
(* Type checking of lvalue                                   *)
type msf_e = MSF.t * Env.venv

let ty_lval env ((msf, venv) as msf_e : msf_e) x ety : msf_e =
  (* First path the type ety to make it consistant with the variable info *)
  match x with
  | Lnone _ -> msf_e
  | Lvar x ->
      (* TODO assumption: p = e when p is a pointer and e a direct value means p
         points to a new position, where the expression is *)
      (* as opposed to assigning the pointer directly to the given value *)
      (* likewise, assuming x = p means storing the value pointed by p in x *)
      (* what with p = [ x ]? I believe it does not compile *)
      let lp, le =
        match ety with
        | Direct le -> Env.public2 env, le
        | Indirect(lp, le) -> lp, le in
      let xty = if is_ptr (kind_i x)
          then Indirect(lp, le)
          else Direct le in
      let msf = MSF.update msf (L.unloc x) in
      let venv = Env.set_ty env venv x xty in
      msf, venv

  | Lmem(_, _, vi, i) ->
      ensure_public env venv vi i;
        (* programmes are assumed to be safe, thus corruption from memory store
           with [x + i] is speculative only *)
      msf, Env.corruption_speculative env venv (content_ty ety)

  | Laset(_, aa, ws, x, i) ->
      ensure_public_address env venv (L.loc x) x;
      ensure_public env venv (L.loc x) i;
      let le = content_ty ety in
      let venv =
        let l = Env.fresh2 env in
        let xty =
          match Env.get_i venv x with
          | Direct lx -> VlPairs.add_le lx l; VlPairs.add_le le l; Direct l
          | Indirect (lp, lx) -> VlPairs.add_le lx l; VlPairs.add_le le l;
              Indirect (lp, l)
        in
        Env.set_ty env venv x xty in

      if ssafe_test x aa ws i then
        msf, venv
      else (* mispeculation has necessarily occured *)
        msf, Env.corruption_speculative env venv le

  | Lasub(_, _, _, x, i) ->
      ensure_public_address env venv (L.loc x) x;
      ensure_public env venv (L.loc x) i;
      let le = content_ty ety in
      let l = Env.fresh2 env in
      let xty =
        match Env.get_i venv x with
        | Direct lx -> VlPairs.add_le lx l; VlPairs.add_le le l; Direct l
        | Indirect (lp, lx) -> VlPairs.add_le lx l; VlPairs.add_le le l;
            Indirect (lp, l)
      in
      msf, Env.set_ty env venv x xty

let ty_lvals1 env (msf_e : msf_e) xs ety : msf_e =
  List.fold_left (fun msf_e x -> ty_lval env msf_e x ety) msf_e xs

let ty_lvals env (msf_e : msf_e) xs tys : msf_e =
  List.fold_left2 (ty_lval env) msf_e xs tys


(* -------------------------------------------------------------- *)
(* declassify                                                     *)
(* TODO ensure declassify cannot occur on potentially corrupted   *)
(* stack values                                                   *)

let sdeclassify = "declassify"

let is_declassify annot =
  Annot.ensure_uniq1 sdeclassify Annot.none annot <> None

let declassify_lvl env (_, s) = (Env.public env, s)

let declassify env = function
  | Direct le          -> Direct (declassify_lvl env le)
  | Indirect (lp, le)  -> Indirect (lp, declassify_lvl env le)

let declassify_ty env annot ty = if is_declassify annot
  then declassify env ty
  else ty

let declassify_tys env annot tys = if is_declassify annot
  then List.map (declassify env) tys
  else tys


(* right now only used by syscall, which only consists of randombytes
   it is thus tailored for this specific function. *)
let ensure_public_address_expr env venv loc e =
  let ety = ty_expr env venv loc e in
  match ety with
  | Direct _ -> ()
  | Indirect (le, _) -> try VlPairs.add_le le (Env.public2 env)
      with Lvl.Unsat unsat -> error_unsat loc unsat pp_expr e ety (Direct (Env.public2 env))

(* --------------------------------------------------------------- *)
let move_msf ~loc env (msf, venv) mso msi =
  let mso = reg_lval ~direct:true loc mso
  and msi = reg_expr ~direct:true loc msi in
  MSF.check_msf msf msi;
  MSF.add mso msf, Env.set_ty env venv mso (Env.dpublic env)

(* --------------------------------------------------------------- *)
(* [ty_instr env msf i] return msf' such that env, msf |- i : msf' *)

let rec ty_instr is_ct_asm fenv env (msf, venv) i =
  let venv = Env.venv_forget (fst i.i_info) venv in
  let msf, venv = ty_instr_r is_ct_asm fenv env (msf, venv) i in
  msf, Env.venv_forget (snd i.i_info) venv

and ty_instr_r is_ct_asm fenv env ((msf,venv) as msf_e :msf_e) i =
  let loc = i.i_loc.L.base_loc in
  match i.i_desc with
  | Csyscall (xs, o, es) ->
    (* TODO: generalize to other syscalls *)
    assert (match o with Syscall_t.RandomBytes _ -> true);
    List.iter (ensure_public_address_expr env venv loc) es;
    (* We don't known what happen to MSF after external function call *)
    ty_lvals1 env (MSF.toinit, venv) xs (Env.dsecret env)

  | Cassgn(mso, _, _, (Pvar x as msi)) when MSF.is_msf msf x.gv ->
    move_msf ~loc env msf_e mso msi

  | Cassgn(x, _, _, e) ->
    let ety = ty_expr env venv loc e in
    ty_lval env msf_e x (declassify_ty env i.i_annot ety)

  | Copn(xs, _, o, es) ->
    begin match is_special o, xs, es with
    | Init_msf, [ms], _ ->
      let ms = reg_lval_opt ~direct:true loc ms in
      let venv = Env.set_init_msf env venv ms in
      let ms = Option.map_default MSF.exact1 MSF.toinit ms in
      ms, venv

    | Init_msf, _, _ -> assert false

    | Update_msf, [mso], [b; msi] ->
      let mso = reg_lval ~direct:true loc mso and msi = reg_expr ~direct:true loc msi in
      (* do not check b, if check_msf_trans succeed then b is public *)
      MSF.check_msf_trans msf msi b;
      let _, venv = ty_lvals1 env (msf, venv) xs (Env.dpublic env) in
      MSF.exact1 mso, venv

    | Update_msf, _, _ -> assert false

    | Mov_msf, [mso], [msi] ->
      move_msf ~loc env msf_e mso msi

    | Mov_msf, _, _ -> assert false

    | Protect, [x], [e; ms] ->
      let _ = reg_lval ~direct:false loc x and _ = reg_expr ~direct:false loc e and
          ms = reg_expr ~direct:true loc ms in
      MSF.check_msf_exact msf ms;
      let xty =
        match ty_expr env venv loc e with
        | Direct (n, _) -> Direct (n, n)
        | Indirect ((n, _), le) -> Indirect ((n, n), le) in

      ty_lval env msf_e x xty

    | Protect, _, _ -> assert false

    | Other, _, _  ->
      let public = not (CT.is_ct_sopn is_ct_asm o) in
      let ety = ty_exprs_max ~public env venv loc es in
      ty_lvals1 env msf_e xs (declassify_ty env i.i_annot ety)
    end

  | Cif(e, c1, c2) ->
    let msf1, msf2 =
      if is_inline i then
        msf, msf
      else begin
          ensure_public env venv loc e;
          MSF.enter_if msf e, MSF.enter_if msf (Papp1(Onot, e))
        end
    in
    let msf1, venv1 = ty_cmd is_ct_asm fenv env (msf1, venv) c1 in
    let msf2, venv2 = ty_cmd is_ct_asm fenv env (msf2, venv) c2 in
    let venv1 = Env.venv_forget (snd i.i_info) venv1 in
    let venv2 = Env.venv_forget (snd i.i_info) venv2 in
    MSF.max msf1 msf2, Env.max env venv1 venv2

  | Cfor(x, (_, e1, e2), c) ->
      ensure_public env venv loc e1;
      ensure_public env venv loc e2;

      (* Live set after the loop guard *)
      let live_at_c =
        let live_after_i = snd i.i_info in
        match c with i :: _ -> Sv.union live_after_i (fst i.i_info) | [] -> live_after_i in

      let msf = MSF.loop env i.i_loc msf in
      (* let w, _ = written_vars [i] in *)
      let venv1 = Env.freshen env venv in (* venv <= venv1 *)
      let msf_e = ty_lval env (msf, venv1) (Lvar x) (Env.dpublic env) in
      let (msf', venv') = ty_cmd is_ct_asm fenv env msf_e c in
      let msf' = MSF.end_loop loc msf msf' in
      let venv1 = Env.venv_forget live_at_c venv1 in
      let venv' = Env.venv_forget live_at_c venv' in
      Env.ensure_le loc venv' venv1; (* venv' <= venv1 *)
      msf', venv1

  | Cwhile(_, c1, e, (_, (_, live_at_c2)), c2) ->
    (* c1; while e do (c2; c1) *)
    (* env, msf <= env1, msf1
       env1, msf1 |- c1 : msf2, env2   env2 |- e : public
       env2, enter_if e msf2 |- c2 : env1, msf1
       --------------------------------------------------------------------------------
       env, msf |- while c1 e c2 : enter_if e msf1
     *)
    let msf1 = MSF.loop env i.i_loc msf in
    (* let w, _ = written_vars [i] in *)
    (* NOTE cannot restrict refreshed variables to local modified vars
       because of memory corruption: if loop body corrupts some stack variable
       constrained to public, the test fails, while marking all stack variables
       as secret is sufficient *)

    let venv1 = Env.freshen env venv in (* venv <= venv1 *)
    let (msf2, venv2) = ty_cmd is_ct_asm fenv env (msf1, venv1) c1 in
    ensure_public env venv2 loc e;
    let (msf', venv') = ty_cmd is_ct_asm fenv env (MSF.enter_if msf2 e, Env.venv_forget live_at_c2 venv2) c2 in
    let _ = MSF.end_loop loc msf1 msf' in
    Env.ensure_le loc venv' venv1; (* venv' <= venv1 *)
    MSF.enter_if msf2 (Papp1(Onot, e)), venv2

  | Ccall (xs, f, es) ->
    let fty = FEnv.get_fty fenv f in
    let modmsf = fty.modmsf in
    let tyout, tyin, resulting_corruption = Env.clone_for_call env fty in

    let input_ty e vfty =
      match vfty with
      | IsMsf ->
        (* we don't check that e is public, it is ensured by being msf *)
        let ms = reg_expr ~direct:true loc e in
        MSF.check_msf_exact msf ms
      | IsNormal ety' ->
        let ety = ty_expr env venv loc e in
        try match ety, ety' with
        | Direct le, Direct le' -> VlPairs.add_le le le'
        | Direct le, Indirect (_, le') -> VlPairs.add_le le le'
        | Indirect(lp, le), Direct le' -> VlPairs.add_le lp (Env.public2 env); VlPairs.add_le le le'
        | Indirect(lp, le), Indirect(lp', le') -> VlPairs.add_le lp lp'; VlPairs.add_le le le'
        with Lvl.Unsat unsat -> error_unsat loc unsat pp_expr e ety ety' in
    List.iter2 input_ty es tyin;

    (* callee function has its own effect on this function corruption *)
    let venv = Env.corruption env venv resulting_corruption in

    (* compute the resulting venv *)
    let output_ty msf_e x vfty =
      let ty =
        match vfty with
        | IsMsf -> Env.dpublic env
        | IsNormal ty -> declassify_ty env i.i_annot ty in
      let (msf, venv) = ty_lval env msf_e x ty in
      let msf = if vfty = IsMsf then MSF.add (reg_lval ~direct:true loc x) msf else msf in
      (msf, venv) in
    let msf = if is_Modified modmsf then MSF.toinit else msf in
    List.fold_left2 output_ty (msf, venv) xs tyout

and ty_cmd is_ct_asm fenv env msf_e c =
  List.fold_left (ty_instr is_ct_asm fenv env) msf_e c


(* ------------------------------------------------------------------- *)
(* Do the inference + type checking of function                        *)
(*
#nomodmsf #constraints = "l1 <= transient, l2 <= l1"
   fn f (#public #secret reg u64[1] x, #poly = l1 stack u8 c) ->
        #poly=l1 #poly=l2 u64[1]
*)

let parse_var_annot ~(msf:bool) (annot: annotations) : ulevel list =
  let module A = Annot in
  let filters =
    [spublic, (fun a -> A.none a; Public);
     ssecret, (fun a -> A.none a; Secret);
     stransient, (fun a -> A.none a; Transient)
     ] in

  let filters =
    if msf then (smsf, (fun a -> A.none a; Msf)) :: filters else filters in

  let lvls = A.process_annot filters annot in

  List.map snd lvls

exception Error_after of string * string

let parse_constraints (s:string) : (string * string) list =
  let error expected s n =
    raise (Error_after (expected, String.sub s 0 n)) in

  let is_blank c = c = ' ' || c = '\n' || c = '\t' in

  let rec check_blank expected s first last =
    if first >= last then ()
    else if is_blank s.[first] then
      check_blank expected s (first + 1) last
    else error expected s first in

  let search expected re s n =
    let first =
      try Str.search_forward re s n
      with Not_found -> error expected s n in
    let sub = Str.matched_string s in
    let next = Str.match_end () in
    check_blank expected s n first;
    sub, next in

  let rec skip_blank s n =
    if n < String.length s && is_blank s.[n] then skip_blank s (n+1) else n in

  let ident s n = search "ident" (Str.regexp "[A-Za-z0-9]+") s n in

  let is_le s n = snd (search "<=" (Str.regexp "<=") s n) in

  let is_comma s n = snd (search "," (Str.regexp ",") s n) in

  let parse_c s n =
    let l1, n = ident s n in
    let n = is_le s n in
    let l2, n = ident s n in
    (l1, l2), n in

  let rec parse_rec acc s n =
    let n = skip_blank s n in
    if String.length s <= n then acc
    else
      let n = is_comma s n in
      let c, n = parse_c s n in
      parse_rec (c::acc) s n in

  let n = skip_blank s 0 in
  if String.length s <= n then []
  else
    let c, n = parse_c s n in
    parse_rec [c] s n

let parse_user_constraints (a:annotations) : (string * string) list =
  let module A = Annot in
  let sconstraints = "constraints" in
  let on_string loc _ s =
    try parse_constraints s
    with Error_after(kw, s) ->
      A.error ~loc
        "%s expected after %s" kw s in
  let error loc =
    A.error ~loc
      "attribute for %s should be a string" sconstraints in

   List.flatten
     (List.map snd (A.process_annot [sconstraints, A.on_attribute ~on_string error] a))

let init_constraint fenv f =
  let sig_annot = SecurityAnnotations.get_sct_signature f.f_annot.f_user_annot in
  let env = Env.init () in
  let venv = Env.empty env in
  let tbl = Hashtbl.create 97 in
  Hashtbl.add tbl spublic    (Env.public2 env);
  Hashtbl.add tbl stransient (Env.transient env);
  Hashtbl.add tbl ssecret    (Env.secret2 env);
  (* export function: all input type should be at most transient and msf is not allowed *)
  (* the new version does not take that into accout, does it? *)
  let export = FInfo.is_export f.f_cc in

  let add_lvl s =
    try Hashtbl.find tbl s
    with Not_found ->
      let l = Env.fresh2 ~name:s env in
      Hashtbl.add tbl s l;
      l in

  let to_lvl = function
    | Secret -> Env.secret2 env
    | Transient -> Env.transient env
    | Public | Msf -> Env.public2 env in

  let lvl_of_sa_level =
    let open SecurityAnnotations in
    let lvl_of_simple_level get =
      function
      | Public -> Env.public env
      | Secret -> Env.secret env
      | Named s -> get (add_lvl s)
    in
    function { normal; speculative } ->
      lvl_of_simple_level fst normal, lvl_of_simple_level snd speculative
  in

  let to_vty =
    function
    | SecurityAnnotations.Msf -> Direct (to_lvl Msf)
    | Direct n -> Direct (lvl_of_sa_level n)
    | Indirect { ptr; value } -> Indirect (lvl_of_sa_level ptr, lvl_of_sa_level value)
  in

  let error_msf loc =
    error ~loc
      "%s annotation not allowed here" smsf in

  (** The [is_local] argument is true when variable [x] is a local variable as
  opposed to an argument or a returned value which inherits constraints from the call-sites. *)
  let mk_vty loc ~is_local ~(msf:bool) x ls an =
    let msf, ovty =
      match ls, an with
      | [], None -> None, None
      | [l], None ->
        if not msf && l = Msf then error_msf loc;
        Some (l = Msf), Some(Direct (to_lvl l))
      | [l1; l2], None ->
        if (l1 = Msf || l2 = Msf) then error_msf loc;
        Some false, Some(Indirect (to_lvl l1, to_lvl l2))
      | _, None ->
        error ~loc:(x.v_dloc)
          "invalid security annotations %a" pp_var x
      | [Msf], Some n ->
         if not msf then error_msf loc;
         Some true, Some (to_vty n)
      | _, Some n ->
         Some (n = SecurityAnnotations.Msf), Some (to_vty n)
    in
    let vty =
      match ovty with
      | None ->
        begin match x.v_kind with
        | Const -> Env.dpublic env
        | (Stack Direct | Reg (_, Direct)) when is_local -> Direct (Env.secret2 env)
        | (Stack Direct | Reg (_, Direct)) -> Direct (Env.fresh2 env)
        | (Stack (Pointer _) | Reg (_, Pointer _)) when is_local -> Indirect(Env.secret2 env, Env.fresh2 env)
        | (Stack (Pointer _) | Reg (_, Pointer _)) -> Indirect(Env.fresh2 env, Env.fresh2 env)
        | Inline -> Env.dpublic env
        | Global -> Env.dpublic env (* unsure *)
        end
      | Some ty -> (* this partly has the same role as Env.init_ty. Remove one occurence? *)
        begin match x.v_kind, ty with
        | Const, Direct _ -> ()
        | Stack Direct, Direct _ -> ()
        | Stack (Pointer _), Indirect _ -> ()
        | Reg (_, Direct), Direct _ -> ()
        | Reg (_, Pointer _), Indirect _ -> ()
        | Inline, Direct _ -> ()
        | Global, Direct _ -> ()
        | _ ->
          error ~loc
            "invalid security annotations for %a" pp_var x
        end; ty in
    msf, vty in

  let process_return i x annot =
    let loc = L.loc x and x = L.unloc x in
    let an = Option.bind sig_annot (SecurityAnnotations.get_nth_result i) in
    let ls = parse_var_annot ~msf:(not export) annot in
    mk_vty ~is_local:false loc ~msf:(not export) x ls an in

  (* process function outputs *)
  let tyout = List.map2i process_return f.f_ret f.f_ret_info.ret_annot in

  (* infer msf_oracle info *)
  let msfs =
    infer_msf_c ~withcheck:true fenv (Env.get_msf_oracle env) f.f_body
      (List.fold_left2 (fun s x (msf, _) ->
           if Option.default false msf then Sv.add (L.unloc x) s else s)
         Sv.empty f.f_ret tyout) in

  if export && not (Sv.is_empty msfs) then begin
    let vars_kind, pos =
      if Sv.subset msfs (Sv.of_list f.f_args)
      then "arguments", ", this is not allowed for export functions"
      else "variables", ""
    in
    error
      ~loc:f.f_loc
      "@[<h>the %s %a need to be MSFs%s.@]"
      vars_kind
      pp_vset msfs
      pos
  end;

  (* process function inputs *)
  let process_param i venv x =
    let an = Option.bind sig_annot (SecurityAnnotations.get_nth_argument i) in
    let ls = parse_var_annot ~msf:(not export) x.v_annot in
    let msf, vty = mk_vty ~is_local:false x.v_dloc ~msf:(not export) x ls an in
    let msf =
      match msf with
      | None -> Sv.mem x msfs
      | Some b ->
        if b <> Sv.mem x msfs then begin
          let loc = x.v_dloc in
          if b
          then warn ~loc:loc "%a does not need to be an MSF" pp_var x
          else error ~loc "%a should be an MSF" pp_var x
        end;
        b in
    if export then
      begin let lvls = match vty with
      | Indirect (p, v) -> [ p; v ]
      | Direct v -> [ v ]
      in List.iter begin fun l ->
          try VlPairs.add_le (Env.public env, Env.secret env) l
          with Lvl.Unsat _unsat ->
            error ~loc:(x.v_dloc)
              "security annotation for %a should be at least %s"
                 pp_var x stransient
        end lvls
      end;
    let venv = Env.add_var env venv x vty in
    let ty = if msf then IsMsf else IsNormal vty in
    venv, ty in

  let venv, tyin = List.mapi_fold process_param venv f.f_args in

  (* build the constraints *)
  let do_constraint (s1, s2) =
    let get s =
      try Hashtbl.find tbl s
      with Not_found ->
        error ~loc:f.f_loc "unbound security level %s" s in
    try VlPairs.add_le (get s1) (get s2)
    with Lvl.Unsat _unsat ->
      error ~loc:f.f_loc
        "cannot add constraint %s <= %s, it is inconsistent" s1 s2 in
  List.iter do_constraint (parse_user_constraints f.f_annot.f_user_annot);

  (* init type for local *)
  let do_local x venv =
    let _, vty = mk_vty x.v_dloc ~is_local:true ~msf:false x [] None in
    Env.add_var env venv x vty in

  let venv = Sv.fold do_local (locals f) venv in

  (* infer modmsf and check consistency with user info *)
  let modmsf = modmsf_c fenv f.f_body in
  let umodmsf =
    Annot.ensure_uniq
       ["modmsf", (fun a -> Annot.none a; true);
        "nomodmsf", (fun a -> Annot.none a; false)] f.f_annot.f_user_annot in
  begin match umodmsf with
  | None -> ()
  | Some annot ->
      if annot <> is_Modified modmsf then
        let sannot = if annot then "modmsf" else "nomodmsf" in
        error ~loc:f.f_loc "annotation %s should be %a" sannot pp_modmsf modmsf
  end;

  env, venv, tyin, tyout, modmsf


let rec ty_fun is_ct_asm fenv fn =
  try Hf.find fenv.env_ty fn
  with Not_found ->
    let fty = ty_fun_infer is_ct_asm fenv fn in
    Hf.add fenv.env_ty fn fty;
    fty

and ty_fun_infer is_ct_asm fenv fn =
  let f = FEnv.get_fun_def fenv fn in
  (* First compute all function call by f and recurse *)
  let _, called = written_vars_fc f in
  Mf.iter (fun fn _ -> ignore (ty_fun is_ct_asm fenv fn)) called;
  let env, venv, tyin, tyout, modmsf = init_constraint fenv f in
  (* init msf status *)
  let msf =
    List.fold_left2 (fun msf x ty ->
        if ty = IsMsf then MSF.add (L.mk_loc x.v_dloc x) msf else msf)
      MSF.toinit f.f_args tyin in
  (* start type checking of the body *)
  let msf, venv = ty_cmd is_ct_asm fenv env (msf, venv) f.f_body in
  (* build the resulting type *)
  let doout x (omsf, ty) =
    let le_ty ty1 ty2 =
      try
        match ty1, ty2 with
        | Direct le1, Direct le2 -> VlPairs.add_le le1 le2
        | Indirect(lp1, le1), Indirect(lp2, le2) ->
          VlPairs.add_le lp1 lp2; VlPairs.add_le le1 le2
        | _, _ -> assert false
      with Lvl.Unsat _unsat ->
        error ~loc:(L.loc x)
          "return type for %a is %a it should be less than %a"
             pp_var_i x pp_vty ty1 pp_vty ty2 in
    match omsf with
    | Some true  -> MSF.check_msf_exact msf x; IsMsf
    | Some false ->
      if MSF.is_msf_exact msf x then
        error ~loc:(L.loc x)
          "return annotation for %a should be %s" pp_var_i x smsf;
      le_ty (Env.get_i venv x) ty;
      IsNormal ty
    | None ->
      if MSF.is_msf_exact msf x then IsMsf
      else (le_ty (Env.get_i venv x) ty; IsNormal ty) in

  let tyout = List.map2 doout f.f_ret tyout in
  let resulting_corruption = Env.get_resulting_corruption venv in
  let (n1, s1) = resulting_corruption in

  let constraints = Env.constraints env in
  let add ls vty =
    match vty with
    | IsMsf -> ls
    | IsNormal (Direct (n, s)) -> n :: s :: ls
    | IsNormal (Indirect ((np, sp), (ne, se))) -> np :: sp :: ne :: se :: ls in
  let to_keep = List.fold_left add (List.fold_left add [n1; s1] tyin) tyout in

  C.prune constraints to_keep;
  let fty = { modmsf; tyin; tyout; constraints; resulting_corruption; } in
  if !Glob_options.debug then
    Format.eprintf
      "Before optimization:@.%a@."
      pp_funty
      (f.f_name.fn_name, fty);
  let tomax = List.fold_left add [] tyin in
  let tomin = List.fold_left add [n1; s1] tyout in
  C.optimize constraints ~tomin ~tomax;
  fty


let ty_prog is_ct_asm (prog:('info, 'asm) prog) fl =
  let prog = Liveness.liveness false prog in
  let prog = snd prog in
  let fenv = { env_ty = Hf.create 101; env_def = prog } in
  let fl =
    if fl = [] then
      List.rev_map (fun f -> f.f_name) prog
    else
      let get fn =
        try (List.find (fun f -> f.f_name.fn_name = fn) prog).f_name
        with Not_found ->
          hierror ~loc:Lnone ~kind:"speculative constant type checker" "unknown function %s" fn in
      List.map get fl in
  List.map (fun fn -> fn.fn_name, ty_fun is_ct_asm fenv fn) fl

(* ------------------------------------------------------------------------------- *)
(* Inference of msf_info needed by the compiler                                    *)

let compile_infer_msf (prog:('info, 'asm) prog) =
  let prog = snd prog in
  let fenv = { env_ty = Hf.create 101; env_def = prog } in

  let env = Env.init () in
  (* dummy infos *)
  let resulting_corruption = Env.public2 env in
  let notmsf = IsNormal (Direct resulting_corruption) in
  let constraints = C.init() in

  let infer_fun f =
    let sig_annot = SecurityAnnotations.get_sct_signature f.f_annot.f_user_annot in

    let process_return i annot =
      let ls = parse_var_annot ~msf:true annot in
      let an = Option.bind sig_annot (SecurityAnnotations.get_nth_result i) in
      List.mem Msf ls || an = Some SecurityAnnotations.Msf
    in

    (* process function outputs *)
    let tyout = List.mapi process_return f.f_ret_info.ret_annot in

    (* infer the set of input variables that need to be msf *)
    let msfin =
      infer_msf_c ~withcheck:false fenv (Hashtbl.create 13) f.f_body
        (List.fold_left2 (fun s x msf ->
           if msf then Sv.add (L.unloc x) s else s)
         Sv.empty f.f_ret tyout)
    in
    let mkmsf msf = if msf then IsMsf else notmsf in
    let tyout = List.map mkmsf tyout in
    let tyin  = List.map (fun x -> mkmsf (Sv.mem x msfin)) f.f_args in

   (* infer modmsf and check consistency with user info *)
   let modmsf = modmsf_c fenv f.f_body in

   let fty = {
       modmsf;
       tyin;
       tyout;
       constraints; (* dummy info *)
       resulting_corruption; (* dummy info *)
     }  in
   Hf.add fenv.env_ty f.f_name fty
  in
  List.iter infer_fun (List.rev prog);

  let do_t = function
     | IsNormal _ -> Slh_lowering.Slh_None
     | IsMsf      -> Slh_msf
  in

  let do_f _fn fty =
    let in_t = List.map do_t fty.tyin in
    let out_t = List.map do_t fty.tyout in
    (in_t, out_t)
  in

  Hf.map do_f fenv.env_ty